add read me
This commit is contained in:
21
venv/lib/python3.12/site-packages/onnxruntime/LICENSE
Normal file
21
venv/lib/python3.12/site-packages/onnxruntime/LICENSE
Normal file
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
21
venv/lib/python3.12/site-packages/onnxruntime/Privacy.md
Normal file
21
venv/lib/python3.12/site-packages/onnxruntime/Privacy.md
Normal file
@@ -0,0 +1,21 @@
|
||||
# Privacy
|
||||
|
||||
## Data Collection
|
||||
The software may collect information about you and your use of the software and send it to Microsoft. Microsoft may use this information to provide services and improve our products and services. You may turn off the telemetry as described in the repository. There are also some features in the software that may enable you and Microsoft to collect data from users of your applications. If you use these features, you must comply with applicable law, including providing appropriate notices to users of your applications together with a copy of Microsoft's privacy statement. Our privacy statement is located at https://go.microsoft.com/fwlink/?LinkID=824704. You can learn more about data collection and use in the help documentation and our privacy statement. Your use of the software operates as your consent to these practices.
|
||||
|
||||
***
|
||||
|
||||
### Private Builds
|
||||
No data collection is performed when using your private builds built from source code.
|
||||
|
||||
### Official Builds
|
||||
ONNX Runtime does not maintain any independent telemetry collection mechanisms outside of what is provided by the platforms it supports. However, where applicable, ONNX Runtime will take advantage of platform-supported telemetry systems to collect trace events with the goal of improving product quality.
|
||||
|
||||
Currently telemetry is only implemented for Windows builds and is turned **ON** by default in the official builds distributed in their respective package management repositories ([see here](../README.md#binaries)). This may be expanded to cover other platforms in the future. Data collection is implemented via 'Platform Telemetry' per vendor platform providers (see [telemetry.h](../onnxruntime/core/platform/telemetry.h)).
|
||||
|
||||
#### Technical Details
|
||||
The Windows provider uses the [TraceLogging](https://docs.microsoft.com/en-us/windows/win32/tracelogging/trace-logging-about) API for its implementation. This enables ONNX Runtime trace events to be collected by the operating system, and based on user consent, this data may be periodically sent to Microsoft servers following GDPR and privacy regulations for anonymity and data access controls.
|
||||
|
||||
Windows ML and onnxruntime C APIs allow Trace Logging to be turned on/off (see [API pages](../README.md#api-documentation) for details).
|
||||
For information on how to enable and disable telemetry, see [C API: Telemetry](./C_API.md#telemetry).
|
||||
There are equivalent APIs in the C#, Python, and Java language bindings as well.
|
||||
6156
venv/lib/python3.12/site-packages/onnxruntime/ThirdPartyNotices.txt
Normal file
6156
venv/lib/python3.12/site-packages/onnxruntime/ThirdPartyNotices.txt
Normal file
File diff suppressed because it is too large
Load Diff
360
venv/lib/python3.12/site-packages/onnxruntime/__init__.py
Normal file
360
venv/lib/python3.12/site-packages/onnxruntime/__init__.py
Normal file
@@ -0,0 +1,360 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
"""
|
||||
ONNX Runtime is a performance-focused scoring engine for Open Neural Network Exchange (ONNX) models.
|
||||
For more information on ONNX Runtime, please see `aka.ms/onnxruntime <https://aka.ms/onnxruntime/>`_
|
||||
or the `Github project <https://github.com/microsoft/onnxruntime/>`_.
|
||||
"""
|
||||
|
||||
__version__ = "1.23.2"
|
||||
__author__ = "Microsoft"
|
||||
|
||||
# we need to do device version validation (for example to check Cuda version for an onnxruntime-training package).
|
||||
# in order to know whether the onnxruntime package is for training it needs
|
||||
# to do import onnxruntime.training.ortmodule first.
|
||||
# onnxruntime.capi._pybind_state is required before import onnxruntime.training.ortmodule.
|
||||
# however, import onnxruntime.capi._pybind_state will already raise an exception if a required Cuda version
|
||||
# is not found.
|
||||
# here we need to save the exception and continue with Cuda version validation in order to post
|
||||
# meaningful messages to the user.
|
||||
# the saved exception is raised after device version validation.
|
||||
try:
|
||||
from onnxruntime.capi._pybind_state import (
|
||||
ExecutionMode, # noqa: F401
|
||||
ExecutionOrder, # noqa: F401
|
||||
GraphOptimizationLevel, # noqa: F401
|
||||
LoraAdapter, # noqa: F401
|
||||
ModelMetadata, # noqa: F401
|
||||
NodeArg, # noqa: F401
|
||||
OrtAllocatorType, # noqa: F401
|
||||
OrtArenaCfg, # noqa: F401
|
||||
OrtCompileApiFlags, # noqa: F401
|
||||
OrtDeviceMemoryType, # noqa: F401
|
||||
OrtEpDevice, # noqa: F401
|
||||
OrtExecutionProviderDevicePolicy, # noqa: F401
|
||||
OrtExternalInitializerInfo, # noqa: F401
|
||||
OrtHardwareDevice, # noqa: F401
|
||||
OrtHardwareDeviceType, # noqa: F401
|
||||
OrtMemoryInfo, # noqa: F401
|
||||
OrtMemoryInfoDeviceType, # noqa: F401
|
||||
OrtMemType, # noqa: F401
|
||||
OrtSparseFormat, # noqa: F401
|
||||
OrtSyncStream, # noqa: F401
|
||||
RunOptions, # noqa: F401
|
||||
SessionIOBinding, # noqa: F401
|
||||
SessionOptions, # noqa: F401
|
||||
create_and_register_allocator, # noqa: F401
|
||||
create_and_register_allocator_v2, # noqa: F401
|
||||
disable_telemetry_events, # noqa: F401
|
||||
enable_telemetry_events, # noqa: F401
|
||||
get_all_providers, # noqa: F401
|
||||
get_available_providers, # noqa: F401
|
||||
get_build_info, # noqa: F401
|
||||
get_device, # noqa: F401
|
||||
get_ep_devices, # noqa: F401
|
||||
get_version_string, # noqa: F401
|
||||
has_collective_ops, # noqa: F401
|
||||
register_execution_provider_library, # noqa: F401
|
||||
set_default_logger_severity, # noqa: F401
|
||||
set_default_logger_verbosity, # noqa: F401
|
||||
set_global_thread_pool_sizes, # noqa: F401
|
||||
set_seed, # noqa: F401
|
||||
unregister_execution_provider_library, # noqa: F401
|
||||
)
|
||||
|
||||
import_capi_exception = None
|
||||
except Exception as e:
|
||||
import_capi_exception = e
|
||||
|
||||
from onnxruntime.capi import onnxruntime_validation
|
||||
|
||||
if import_capi_exception:
|
||||
raise import_capi_exception
|
||||
|
||||
from onnxruntime.capi.onnxruntime_inference_collection import (
|
||||
AdapterFormat, # noqa: F401
|
||||
InferenceSession, # noqa: F401
|
||||
IOBinding, # noqa: F401
|
||||
ModelCompiler, # noqa: F401
|
||||
OrtDevice, # noqa: F401
|
||||
OrtValue, # noqa: F401
|
||||
SparseTensor, # noqa: F401
|
||||
copy_tensors, # noqa: F401
|
||||
)
|
||||
|
||||
# TODO: thiagofc: Temporary experimental namespace for new PyTorch front-end
|
||||
try: # noqa: SIM105
|
||||
from . import experimental # noqa: F401
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
package_name, version, cuda_version = onnxruntime_validation.get_package_name_and_version_info()
|
||||
|
||||
if version:
|
||||
__version__ = version
|
||||
|
||||
onnxruntime_validation.check_distro_info()
|
||||
|
||||
|
||||
def _get_package_version(package_name: str):
|
||||
from importlib.metadata import PackageNotFoundError, version # noqa: PLC0415
|
||||
|
||||
try:
|
||||
package_version = version(package_name)
|
||||
except PackageNotFoundError:
|
||||
package_version = None
|
||||
return package_version
|
||||
|
||||
|
||||
def _get_package_root(package_name: str, directory_name: str | None = None):
|
||||
from importlib.metadata import PackageNotFoundError, distribution # noqa: PLC0415
|
||||
|
||||
root_directory_name = directory_name or package_name
|
||||
try:
|
||||
dist = distribution(package_name)
|
||||
files = dist.files or []
|
||||
|
||||
for file in files:
|
||||
if file.name.endswith("__init__.py") and root_directory_name in file.parts:
|
||||
return file.locate().parent
|
||||
|
||||
# Fallback to the first __init__.py
|
||||
if not directory_name:
|
||||
for file in files:
|
||||
if file.name.endswith("__init__.py"):
|
||||
return file.locate().parent
|
||||
except PackageNotFoundError:
|
||||
# package not found, do nothing
|
||||
pass
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _get_nvidia_dll_paths(is_windows: bool, cuda: bool = True, cudnn: bool = True):
|
||||
if is_windows:
|
||||
# Path is relative to site-packages directory.
|
||||
cuda_dll_paths = [
|
||||
("nvidia", "cublas", "bin", "cublasLt64_12.dll"),
|
||||
("nvidia", "cublas", "bin", "cublas64_12.dll"),
|
||||
("nvidia", "cufft", "bin", "cufft64_11.dll"),
|
||||
("nvidia", "cuda_runtime", "bin", "cudart64_12.dll"),
|
||||
]
|
||||
cudnn_dll_paths = [
|
||||
("nvidia", "cudnn", "bin", "cudnn_engines_runtime_compiled64_9.dll"),
|
||||
("nvidia", "cudnn", "bin", "cudnn_engines_precompiled64_9.dll"),
|
||||
("nvidia", "cudnn", "bin", "cudnn_heuristic64_9.dll"),
|
||||
("nvidia", "cudnn", "bin", "cudnn_ops64_9.dll"),
|
||||
("nvidia", "cudnn", "bin", "cudnn_adv64_9.dll"),
|
||||
("nvidia", "cudnn", "bin", "cudnn_graph64_9.dll"),
|
||||
("nvidia", "cudnn", "bin", "cudnn64_9.dll"),
|
||||
]
|
||||
else: # Linux
|
||||
# cublas64 depends on cublasLt64, so cublasLt64 should be loaded first.
|
||||
cuda_dll_paths = [
|
||||
("nvidia", "cublas", "lib", "libcublasLt.so.12"),
|
||||
("nvidia", "cublas", "lib", "libcublas.so.12"),
|
||||
("nvidia", "cuda_nvrtc", "lib", "libnvrtc.so.12"),
|
||||
("nvidia", "curand", "lib", "libcurand.so.10"),
|
||||
("nvidia", "cufft", "lib", "libcufft.so.11"),
|
||||
("nvidia", "cuda_runtime", "lib", "libcudart.so.12"),
|
||||
]
|
||||
|
||||
# Do not load cudnn sub DLLs (they will be dynamically loaded later) to be consistent with PyTorch in Linux.
|
||||
cudnn_dll_paths = [
|
||||
("nvidia", "cudnn", "lib", "libcudnn.so.9"),
|
||||
]
|
||||
|
||||
return (cuda_dll_paths if cuda else []) + (cudnn_dll_paths if cudnn else [])
|
||||
|
||||
|
||||
def print_debug_info():
|
||||
"""Print information to help debugging."""
|
||||
import importlib.util # noqa: PLC0415
|
||||
import os # noqa: PLC0415
|
||||
import platform # noqa: PLC0415
|
||||
from importlib.metadata import distributions # noqa: PLC0415
|
||||
|
||||
print(f"{package_name} version: {__version__}")
|
||||
if cuda_version:
|
||||
print(f"CUDA version used in build: {cuda_version}")
|
||||
print("platform:", platform.platform())
|
||||
|
||||
print("\nPython package, version and location:")
|
||||
ort_packages = []
|
||||
for dist in distributions():
|
||||
package = dist.metadata["Name"]
|
||||
if package == "onnxruntime" or package.startswith(("onnxruntime-", "ort-")):
|
||||
# Exclude packages whose root directory name is not onnxruntime.
|
||||
location = _get_package_root(package, "onnxruntime")
|
||||
if location and (package not in ort_packages):
|
||||
ort_packages.append(package)
|
||||
print(f"{package}=={dist.version} at {location}")
|
||||
|
||||
if len(ort_packages) > 1:
|
||||
print(
|
||||
"\033[33mWARNING: multiple onnxruntime packages are installed to the same location. "
|
||||
"Please 'pip uninstall` all above packages, then `pip install` only one of them.\033[0m"
|
||||
)
|
||||
|
||||
if cuda_version:
|
||||
# Print version of installed packages that is related to CUDA or cuDNN DLLs.
|
||||
packages = [
|
||||
"torch",
|
||||
"nvidia-cuda-runtime-cu12",
|
||||
"nvidia-cudnn-cu12",
|
||||
"nvidia-cublas-cu12",
|
||||
"nvidia-cufft-cu12",
|
||||
"nvidia-curand-cu12",
|
||||
"nvidia-cuda-nvrtc-cu12",
|
||||
"nvidia-nvjitlink-cu12",
|
||||
]
|
||||
for package in packages:
|
||||
directory_name = "nvidia" if package.startswith("nvidia-") else None
|
||||
version = _get_package_version(package)
|
||||
if version:
|
||||
print(f"{package}=={version} at {_get_package_root(package, directory_name)}")
|
||||
else:
|
||||
print(f"{package} not installed")
|
||||
|
||||
if platform.system() == "Windows":
|
||||
print(f"\nEnvironment variable:\nPATH={os.environ['PATH']}")
|
||||
elif platform.system() == "Linux":
|
||||
print(f"\nEnvironment variable:\nLD_LIBRARY_PATH={os.environ['LD_LIBRARY_PATH']}")
|
||||
|
||||
if importlib.util.find_spec("psutil"):
|
||||
|
||||
def is_target_dll(path: str):
|
||||
target_keywords = ["vcruntime140", "msvcp140"]
|
||||
if cuda_version:
|
||||
target_keywords = ["cufft", "cublas", "cudart", "nvrtc", "curand", "cudnn", *target_keywords]
|
||||
return any(keyword in path for keyword in target_keywords)
|
||||
|
||||
import psutil # noqa: PLC0415
|
||||
|
||||
p = psutil.Process(os.getpid())
|
||||
|
||||
print("\nList of loaded DLLs:")
|
||||
for lib in p.memory_maps():
|
||||
if is_target_dll(lib.path.lower()):
|
||||
print(lib.path)
|
||||
|
||||
if cuda_version:
|
||||
if importlib.util.find_spec("cpuinfo") and importlib.util.find_spec("py3nvml"):
|
||||
from .transformers.machine_info import get_device_info # noqa: PLC0415
|
||||
|
||||
print("\nDevice information:")
|
||||
print(get_device_info())
|
||||
else:
|
||||
print("please `pip install py-cpuinfo py3nvml` to show device information.")
|
||||
else:
|
||||
print("please `pip install psutil` to show loaded DLLs.")
|
||||
|
||||
|
||||
def preload_dlls(cuda: bool = True, cudnn: bool = True, msvc: bool = True, directory=None):
|
||||
"""Preload CUDA 12.x and cuDNN 9.x DLLs in Windows or Linux, and MSVC runtime DLLs in Windows.
|
||||
|
||||
When the installed PyTorch is compatible (using same major version of CUDA and cuDNN),
|
||||
there is no need to call this function if `import torch` is done before `import onnxruntime`.
|
||||
|
||||
Args:
|
||||
cuda (bool, optional): enable loading CUDA DLLs. Defaults to True.
|
||||
cudnn (bool, optional): enable loading cuDNN DLLs. Defaults to True.
|
||||
msvc (bool, optional): enable loading MSVC DLLs in Windows. Defaults to True.
|
||||
directory(str, optional): a directory contains CUDA or cuDNN DLLs. It can be an absolute path,
|
||||
or a path relative to the directory of this file.
|
||||
If directory is None (default value), the search order: the lib directory of compatible PyTorch in Windows,
|
||||
nvidia site packages, default DLL loading paths.
|
||||
If directory is empty string (""), the search order: nvidia site packages, default DLL loading paths.
|
||||
If directory is a path, the search order: the directory, default DLL loading paths.
|
||||
"""
|
||||
import ctypes # noqa: PLC0415
|
||||
import os # noqa: PLC0415
|
||||
import platform # noqa: PLC0415
|
||||
import sys # noqa: PLC0415
|
||||
|
||||
if platform.system() not in ["Windows", "Linux"]:
|
||||
return
|
||||
|
||||
is_windows = platform.system() == "Windows"
|
||||
if is_windows and msvc:
|
||||
try:
|
||||
ctypes.CDLL("vcruntime140.dll")
|
||||
ctypes.CDLL("msvcp140.dll")
|
||||
if platform.machine() != "ARM64":
|
||||
ctypes.CDLL("vcruntime140_1.dll")
|
||||
except OSError:
|
||||
print("Microsoft Visual C++ Redistributable is not installed, this may lead to the DLL load failure.")
|
||||
print("It can be downloaded at https://aka.ms/vs/17/release/vc_redist.x64.exe.")
|
||||
|
||||
if not (cuda_version and cuda_version.startswith("12.")) and (cuda or cudnn):
|
||||
print(
|
||||
f"\033[33mWARNING: {package_name} is not built with CUDA 12.x support. "
|
||||
"Please install a version that supports CUDA 12.x, or call preload_dlls with cuda=False and cudnn=False.\033[0m"
|
||||
)
|
||||
return
|
||||
|
||||
if not (cuda_version and cuda_version.startswith("12.") and (cuda or cudnn)):
|
||||
return
|
||||
|
||||
is_cuda_cudnn_imported_by_torch = False
|
||||
|
||||
if is_windows:
|
||||
torch_version = _get_package_version("torch")
|
||||
is_torch_for_cuda_12 = torch_version and "+cu12" in torch_version
|
||||
if "torch" in sys.modules:
|
||||
is_cuda_cudnn_imported_by_torch = is_torch_for_cuda_12
|
||||
if (torch_version and "+cu" in torch_version) and not is_torch_for_cuda_12:
|
||||
print(
|
||||
f"\033[33mWARNING: The installed PyTorch {torch_version} does not support CUDA 12.x. "
|
||||
f"Please install PyTorch for CUDA 12.x to be compatible with {package_name}.\033[0m"
|
||||
)
|
||||
|
||||
if is_torch_for_cuda_12 and directory is None:
|
||||
torch_root = _get_package_root("torch", "torch")
|
||||
if torch_root:
|
||||
directory = os.path.join(torch_root, "lib")
|
||||
|
||||
base_directory = directory or ".."
|
||||
if not os.path.isabs(base_directory):
|
||||
base_directory = os.path.join(os.path.dirname(__file__), base_directory)
|
||||
base_directory = os.path.normpath(base_directory)
|
||||
if not os.path.isdir(base_directory):
|
||||
raise RuntimeError(f"Invalid parameter of directory={directory}. The directory does not exist!")
|
||||
|
||||
if is_cuda_cudnn_imported_by_torch:
|
||||
# In Windows, PyTorch has loaded CUDA and cuDNN DLLs during `import torch`, no need to load them again.
|
||||
print("Skip loading CUDA and cuDNN DLLs since torch is imported.")
|
||||
return
|
||||
|
||||
# Try load DLLs from nvidia site packages.
|
||||
dll_paths = _get_nvidia_dll_paths(is_windows, cuda, cudnn)
|
||||
loaded_dlls = []
|
||||
for relative_path in dll_paths:
|
||||
dll_path = (
|
||||
os.path.join(base_directory, relative_path[-1])
|
||||
if directory
|
||||
else os.path.join(base_directory, *relative_path)
|
||||
)
|
||||
if os.path.isfile(dll_path):
|
||||
try:
|
||||
_ = ctypes.CDLL(dll_path)
|
||||
loaded_dlls.append(relative_path[-1])
|
||||
except Exception as e:
|
||||
print(f"Failed to load {dll_path}: {e}")
|
||||
|
||||
# Try load DLLs with default path settings.
|
||||
has_failure = False
|
||||
for relative_path in dll_paths:
|
||||
dll_filename = relative_path[-1]
|
||||
if dll_filename not in loaded_dlls:
|
||||
try:
|
||||
_ = ctypes.CDLL(dll_filename)
|
||||
except Exception as e:
|
||||
has_failure = True
|
||||
print(f"Failed to load {dll_filename}: {e}")
|
||||
|
||||
if has_failure:
|
||||
print("Please follow https://onnxruntime.ai/docs/install/#cuda-and-cudnn to install CUDA and CuDNN.")
|
||||
Binary file not shown.
@@ -0,0 +1,6 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
from .backend import is_compatible, prepare, run, supports_device # noqa: F401
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
175
venv/lib/python3.12/site-packages/onnxruntime/backend/backend.py
Normal file
175
venv/lib/python3.12/site-packages/onnxruntime/backend/backend.py
Normal file
@@ -0,0 +1,175 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
"""
|
||||
Implements ONNX's backend API.
|
||||
"""
|
||||
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import packaging.version
|
||||
from onnx import ModelProto, helper, version # noqa: F401
|
||||
from onnx.backend.base import Backend
|
||||
from onnx.checker import check_model
|
||||
|
||||
from onnxruntime import InferenceSession, SessionOptions, get_available_providers, get_device
|
||||
from onnxruntime.backend.backend_rep import OnnxRuntimeBackendRep
|
||||
|
||||
|
||||
class OnnxRuntimeBackend(Backend):
|
||||
"""
|
||||
Implements
|
||||
`ONNX's backend API <https://github.com/onnx/onnx/blob/main/docs/ImplementingAnOnnxBackend.md>`_
|
||||
with *ONNX Runtime*.
|
||||
The backend is mostly used when you need to switch between
|
||||
multiple runtimes with the same API.
|
||||
`Importing models from ONNX to Caffe2 <https://github.com/onnx/tutorials/blob/master/tutorials/OnnxCaffe2Import.ipynb>`_
|
||||
shows how to use *caffe2* as a backend for a converted model.
|
||||
Note: This is not the official Python API.
|
||||
"""
|
||||
|
||||
allowReleasedOpsetsOnly = bool(os.getenv("ALLOW_RELEASED_ONNX_OPSET_ONLY", "1") == "1") # noqa: N815
|
||||
|
||||
@classmethod
|
||||
def is_compatible(cls, model, device=None, **kwargs):
|
||||
"""
|
||||
Return whether the model is compatible with the backend.
|
||||
|
||||
:param model: unused
|
||||
:param device: None to use the default device or a string (ex: `'CPU'`)
|
||||
:return: boolean
|
||||
"""
|
||||
if device is None:
|
||||
device = get_device()
|
||||
return cls.supports_device(device)
|
||||
|
||||
@classmethod
|
||||
def is_opset_supported(cls, model):
|
||||
"""
|
||||
Return whether the opset for the model is supported by the backend.
|
||||
When By default only released onnx opsets are allowed by the backend
|
||||
To test new opsets env variable ALLOW_RELEASED_ONNX_OPSET_ONLY should be set to 0
|
||||
|
||||
:param model: Model whose opsets needed to be verified.
|
||||
:return: boolean and error message if opset is not supported.
|
||||
"""
|
||||
if cls.allowReleasedOpsetsOnly:
|
||||
for opset in model.opset_import:
|
||||
domain = opset.domain if opset.domain else "ai.onnx"
|
||||
try:
|
||||
key = (domain, opset.version)
|
||||
if key not in helper.OP_SET_ID_VERSION_MAP:
|
||||
error_message = (
|
||||
"Skipping this test as only released onnx opsets are supported."
|
||||
"To run this test set env variable ALLOW_RELEASED_ONNX_OPSET_ONLY to 0."
|
||||
f" Got Domain '{domain}' version '{opset.version}'."
|
||||
)
|
||||
return False, error_message
|
||||
except AttributeError:
|
||||
# for some CI pipelines accessing helper.OP_SET_ID_VERSION_MAP
|
||||
# is generating attribute error. TODO investigate the pipelines to
|
||||
# fix this error. Falling back to a simple version check when this error is encountered
|
||||
if (domain == "ai.onnx" and opset.version > 12) or (domain == "ai.ommx.ml" and opset.version > 2):
|
||||
error_message = (
|
||||
"Skipping this test as only released onnx opsets are supported."
|
||||
"To run this test set env variable ALLOW_RELEASED_ONNX_OPSET_ONLY to 0."
|
||||
f" Got Domain '{domain}' version '{opset.version}'."
|
||||
)
|
||||
return False, error_message
|
||||
return True, ""
|
||||
|
||||
@classmethod
|
||||
def supports_device(cls, device):
|
||||
"""
|
||||
Check whether the backend is compiled with particular device support.
|
||||
In particular it's used in the testing suite.
|
||||
"""
|
||||
if device == "CUDA":
|
||||
device = "GPU"
|
||||
return "-" + device in get_device() or device + "-" in get_device() or device == get_device()
|
||||
|
||||
@classmethod
|
||||
def prepare(cls, model, device=None, **kwargs):
|
||||
"""
|
||||
Load the model and creates a :class:`onnxruntime.InferenceSession`
|
||||
ready to be used as a backend.
|
||||
|
||||
:param model: ModelProto (returned by `onnx.load`),
|
||||
string for a filename or bytes for a serialized model
|
||||
:param device: requested device for the computation,
|
||||
None means the default one which depends on
|
||||
the compilation settings
|
||||
:param kwargs: see :class:`onnxruntime.SessionOptions`
|
||||
:return: :class:`onnxruntime.InferenceSession`
|
||||
"""
|
||||
if isinstance(model, OnnxRuntimeBackendRep):
|
||||
return model
|
||||
elif isinstance(model, InferenceSession):
|
||||
return OnnxRuntimeBackendRep(model)
|
||||
elif isinstance(model, (str, bytes)):
|
||||
options = SessionOptions()
|
||||
for k, v in kwargs.items():
|
||||
if hasattr(options, k):
|
||||
setattr(options, k, v)
|
||||
|
||||
excluded_providers = os.getenv("ORT_ONNX_BACKEND_EXCLUDE_PROVIDERS", default="").split(",")
|
||||
providers = [x for x in get_available_providers() if (x not in excluded_providers)]
|
||||
|
||||
inf = InferenceSession(model, sess_options=options, providers=providers)
|
||||
# backend API is primarily used for ONNX test/validation. As such, we should disable session.run() fallback
|
||||
# which may hide test failures.
|
||||
inf.disable_fallback()
|
||||
if device is not None and not cls.supports_device(device):
|
||||
raise RuntimeError(f"Incompatible device expected '{device}', got '{get_device()}'")
|
||||
return cls.prepare(inf, device, **kwargs)
|
||||
else:
|
||||
# type: ModelProto
|
||||
# check_model serializes the model anyways, so serialize the model once here
|
||||
# and reuse it below in the cls.prepare call to avoid an additional serialization
|
||||
# only works with onnx >= 1.10.0 hence the version check
|
||||
onnx_version = packaging.version.parse(version.version) or packaging.version.Version("0")
|
||||
onnx_supports_serialized_model_check = onnx_version.release >= (1, 10, 0)
|
||||
bin_or_model = model.SerializeToString() if onnx_supports_serialized_model_check else model
|
||||
check_model(bin_or_model)
|
||||
opset_supported, error_message = cls.is_opset_supported(model)
|
||||
if not opset_supported:
|
||||
raise unittest.SkipTest(error_message)
|
||||
# Now bin might be serialized, if it's not we need to serialize it otherwise we'll have
|
||||
# an infinite recursive call
|
||||
bin = bin_or_model
|
||||
if not isinstance(bin, (str, bytes)):
|
||||
bin = bin.SerializeToString()
|
||||
return cls.prepare(bin, device, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def run_model(cls, model, inputs, device=None, **kwargs):
|
||||
"""
|
||||
Compute the prediction.
|
||||
|
||||
:param model: :class:`onnxruntime.InferenceSession` returned
|
||||
by function *prepare*
|
||||
:param inputs: inputs
|
||||
:param device: requested device for the computation,
|
||||
None means the default one which depends on
|
||||
the compilation settings
|
||||
:param kwargs: see :class:`onnxruntime.RunOptions`
|
||||
:return: predictions
|
||||
"""
|
||||
rep = cls.prepare(model, device, **kwargs)
|
||||
return rep.run(inputs, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def run_node(cls, node, inputs, device=None, outputs_info=None, **kwargs):
|
||||
"""
|
||||
This method is not implemented as it is much more efficient
|
||||
to run a whole model than every node independently.
|
||||
"""
|
||||
raise NotImplementedError("It is much more efficient to run a whole model than every node independently.")
|
||||
|
||||
|
||||
is_compatible = OnnxRuntimeBackend.is_compatible
|
||||
prepare = OnnxRuntimeBackend.prepare
|
||||
run = OnnxRuntimeBackend.run_model
|
||||
supports_device = OnnxRuntimeBackend.supports_device
|
||||
@@ -0,0 +1,52 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
"""
|
||||
Implements ONNX's backend API.
|
||||
"""
|
||||
|
||||
from onnx.backend.base import BackendRep
|
||||
|
||||
from onnxruntime import RunOptions
|
||||
|
||||
|
||||
class OnnxRuntimeBackendRep(BackendRep):
|
||||
"""
|
||||
Computes the prediction for a pipeline converted into
|
||||
an :class:`onnxruntime.InferenceSession` node.
|
||||
"""
|
||||
|
||||
def __init__(self, session):
|
||||
"""
|
||||
:param session: :class:`onnxruntime.InferenceSession`
|
||||
"""
|
||||
self._session = session
|
||||
|
||||
def run(self, inputs, **kwargs): # type: (Any, **Any) -> Tuple[Any, ...]
|
||||
"""
|
||||
Computes the prediction.
|
||||
See :meth:`onnxruntime.InferenceSession.run`.
|
||||
"""
|
||||
|
||||
options = RunOptions()
|
||||
for k, v in kwargs.items():
|
||||
if hasattr(options, k):
|
||||
setattr(options, k, v)
|
||||
|
||||
if isinstance(inputs, list):
|
||||
inps = {}
|
||||
for i, inp in enumerate(self._session.get_inputs()):
|
||||
inps[inp.name] = inputs[i]
|
||||
outs = self._session.run(None, inps, options)
|
||||
if isinstance(outs, list):
|
||||
return outs
|
||||
else:
|
||||
output_names = [o.name for o in self._session.get_outputs()]
|
||||
return [outs[name] for name in output_names]
|
||||
else:
|
||||
inp = self._session.get_inputs()
|
||||
if len(inp) != 1:
|
||||
raise RuntimeError(f"Model expect {len(inp)} inputs")
|
||||
inps = {inp[0].name: inputs}
|
||||
return self._session.run(None, inps, options)
|
||||
@@ -0,0 +1,4 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,7 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
# This file can be modified by setup.py when building a manylinux2010 wheel
|
||||
# When modified, it will preload some libraries needed for the python C extension
|
||||
@@ -0,0 +1,33 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
"""
|
||||
Ensure that dependencies are available and then load the extension module.
|
||||
"""
|
||||
import os
|
||||
import platform
|
||||
import warnings
|
||||
|
||||
from . import _ld_preload # noqa: F401
|
||||
|
||||
if platform.system() == "Windows":
|
||||
from . import version_info
|
||||
|
||||
# If on Windows, check if this import error is caused by the user not installing the 2019 VC Runtime
|
||||
# The VC Redist installer usually puts the VC Runtime dlls in the System32 folder, but it may also be found
|
||||
# in some other locations.
|
||||
# TODO, we may want to try to load the VC Runtime dlls instead of checking if the hardcoded file path
|
||||
# is valid, and raise ImportError if the load fails
|
||||
if version_info.vs2019 and platform.architecture()[0] == "64bit":
|
||||
system_root = os.getenv("SystemRoot") or "C:\\Windows"
|
||||
if not os.path.isfile(os.path.join(system_root, "System32", "vcruntime140_1.dll")):
|
||||
warnings.warn("Please install the 2019 Visual C++ runtime and then try again. "
|
||||
"If you've installed the runtime in a non-standard location "
|
||||
"(other than %SystemRoot%\\System32), "
|
||||
"make sure it can be found by setting the correct path.")
|
||||
|
||||
|
||||
|
||||
from .onnxruntime_pybind11_state import * # noqa
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
package_name = 'onnxruntime'
|
||||
__version__ = '1.23.2'
|
||||
@@ -0,0 +1,48 @@
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
|
||||
# This script helps converting .npz files to .onnx_adapter files
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
import onnxruntime as ort
|
||||
|
||||
|
||||
def get_args() -> argparse:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--npz_file_path", type=str, required=True)
|
||||
parser.add_argument("--output_file_path", type=str, required=True)
|
||||
parser.add_argument("--adapter_version", type=int, required=True)
|
||||
parser.add_argument("--model_version", type=int, required=True)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def export_lora_parameters(
|
||||
npz_file_path: os.PathLike, adapter_version: int, model_version: int, output_file_path: os.PathLike
|
||||
):
|
||||
"""The function converts lora parameters in npz to onnx_adapter format"""
|
||||
adapter_format = ort.AdapterFormat()
|
||||
adapter_format.set_adapter_version(adapter_version)
|
||||
adapter_format.set_model_version(model_version)
|
||||
name_to_ort_value = {}
|
||||
with np.load(npz_file_path) as data:
|
||||
for name, np_arr in data.items():
|
||||
ort_value = ort.OrtValue.ortvalue_from_numpy(np_arr)
|
||||
name_to_ort_value[name] = ort_value
|
||||
|
||||
adapter_format.set_parameters(name_to_ort_value)
|
||||
adapter_format.export_adapter(output_file_path)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = get_args()
|
||||
export_lora_parameters(args.npz_file_path, args.adapter_version, args.model_version, args.output_file_path)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
BIN
venv/lib/python3.12/site-packages/onnxruntime/capi/libonnxruntime.so.1.23.2
Executable file
BIN
venv/lib/python3.12/site-packages/onnxruntime/capi/libonnxruntime.so.1.23.2
Executable file
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,47 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
import ctypes
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
|
||||
def find_cudart_versions(build_env=False, build_cuda_version=None):
|
||||
# ctypes.CDLL and ctypes.util.find_library load the latest installed library.
|
||||
# it may not the the library that would be loaded by onnxruntime.
|
||||
# for example, in an environment with Cuda 11.1 and subsequently
|
||||
# conda cudatoolkit 10.2.89 installed. ctypes will find cudart 10.2. however,
|
||||
# onnxruntime built with Cuda 11.1 will find and load cudart for Cuda 11.1.
|
||||
# for the above reason, we need find all versions in the environment and
|
||||
# only give warnings if the expected cuda version is not found.
|
||||
# in onnxruntime build environment, we expected only one Cuda version.
|
||||
if not sys.platform.startswith("linux"):
|
||||
warnings.warn("find_cudart_versions only works on Linux")
|
||||
return None
|
||||
|
||||
cudart_possible_versions = {None, build_cuda_version}
|
||||
|
||||
def get_cudart_version(find_cudart_version=None):
|
||||
cudart_lib_filename = "libcudart.so"
|
||||
if find_cudart_version:
|
||||
cudart_lib_filename = cudart_lib_filename + "." + find_cudart_version
|
||||
|
||||
try:
|
||||
cudart = ctypes.CDLL(cudart_lib_filename)
|
||||
cudart.cudaRuntimeGetVersion.restype = int
|
||||
cudart.cudaRuntimeGetVersion.argtypes = [ctypes.POINTER(ctypes.c_int)]
|
||||
version = ctypes.c_int()
|
||||
status = cudart.cudaRuntimeGetVersion(ctypes.byref(version))
|
||||
if status != 0:
|
||||
return None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
return version.value
|
||||
|
||||
# use set to avoid duplications
|
||||
cudart_found_versions = {get_cudart_version(cudart_version) for cudart_version in cudart_possible_versions}
|
||||
|
||||
# convert to list and remove None
|
||||
return [ver for ver in cudart_found_versions if ver]
|
||||
File diff suppressed because it is too large
Load Diff
Binary file not shown.
@@ -0,0 +1,154 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
"""
|
||||
Check OS requirements for ONNX Runtime Python Bindings.
|
||||
"""
|
||||
|
||||
import linecache
|
||||
import platform
|
||||
import warnings
|
||||
|
||||
|
||||
def check_distro_info():
|
||||
__my_distro__ = ""
|
||||
__my_distro_ver__ = ""
|
||||
__my_system__ = platform.system().lower()
|
||||
|
||||
__OS_RELEASE_FILE__ = "/etc/os-release" # noqa: N806
|
||||
__LSB_RELEASE_FILE__ = "/etc/lsb-release" # noqa: N806
|
||||
|
||||
if __my_system__ == "windows":
|
||||
__my_distro__ = __my_system__
|
||||
__my_distro_ver__ = platform.release().lower()
|
||||
|
||||
if __my_distro_ver__ not in ["10", "11"]:
|
||||
warnings.warn(
|
||||
f"Unsupported Windows version ({__my_distro_ver__}). ONNX Runtime supports Windows 10 and above, only."
|
||||
)
|
||||
elif __my_system__ == "linux":
|
||||
"""Although the 'platform' python module for getting Distro information works well on standard OS images
|
||||
running on real hardware, it is not accurate when running on Azure VMs, Git Bash, Cygwin, etc.
|
||||
The returned values for release and version are unpredictable for virtualized or emulated environments.
|
||||
/etc/os-release and /etc/lsb_release files, on the other hand, are guaranteed to exist and have standard values
|
||||
in all OSes supported by onnxruntime. The former is the current standard file to check OS info and the latter
|
||||
is its predecessor.
|
||||
"""
|
||||
# Newer systems have /etc/os-release with relevant distro info
|
||||
__my_distro__ = linecache.getline(__OS_RELEASE_FILE__, 3)[3:-1]
|
||||
__my_distro_ver__ = linecache.getline(__OS_RELEASE_FILE__, 6)[12:-2]
|
||||
|
||||
# Older systems may have /etc/os-release instead
|
||||
if not __my_distro__:
|
||||
__my_distro__ = linecache.getline(__LSB_RELEASE_FILE__, 1)[11:-1]
|
||||
__my_distro_ver__ = linecache.getline(__LSB_RELEASE_FILE__, 2)[16:-1]
|
||||
|
||||
# Instead of trying to parse distro specific files,
|
||||
# warn the user ONNX Runtime may not work out of the box
|
||||
__my_distro__ = __my_distro__.lower()
|
||||
__my_distro_ver__ = __my_distro_ver__.lower()
|
||||
elif __my_system__ == "darwin":
|
||||
__my_distro__ = __my_system__
|
||||
__my_distro_ver__ = platform.release().lower()
|
||||
|
||||
if int(__my_distro_ver__.split(".")[0]) < 11:
|
||||
warnings.warn(
|
||||
f"Unsupported macOS version ({__my_distro_ver__}). ONNX Runtime supports macOS 11.0 or later."
|
||||
)
|
||||
elif __my_system__ == "aix":
|
||||
import subprocess # noqa: PLC0415
|
||||
|
||||
returned_output = subprocess.check_output("oslevel")
|
||||
__my_distro_ver__str = returned_output.decode("utf-8")
|
||||
__my_distro_ver = __my_distro_ver__str[:3]
|
||||
else:
|
||||
warnings.warn(
|
||||
f"Unsupported platform ({__my_system__}). ONNX Runtime supports Linux, macOS, AIX and Windows platforms, only."
|
||||
)
|
||||
|
||||
|
||||
def get_package_name_and_version_info():
|
||||
package_name = ""
|
||||
version = ""
|
||||
cuda_version = ""
|
||||
|
||||
try:
|
||||
from .build_and_package_info import __version__ as version # noqa: PLC0415
|
||||
from .build_and_package_info import package_name # noqa: PLC0415
|
||||
|
||||
try: # noqa: SIM105
|
||||
from .build_and_package_info import cuda_version # noqa: PLC0415
|
||||
except ImportError:
|
||||
# cuda_version is optional. For example, cpu only package does not have the attribute.
|
||||
pass
|
||||
except Exception as e:
|
||||
warnings.warn("WARNING: failed to collect package name and version info")
|
||||
print(e)
|
||||
|
||||
return package_name, version, cuda_version
|
||||
|
||||
|
||||
def check_training_module():
|
||||
import_ortmodule_exception = None
|
||||
|
||||
has_ortmodule = False
|
||||
try:
|
||||
from onnxruntime.training.ortmodule import ORTModule # noqa: F401, PLC0415
|
||||
|
||||
has_ortmodule = True
|
||||
except ImportError:
|
||||
# ORTModule not present
|
||||
has_ortmodule = False
|
||||
except Exception as e:
|
||||
# this may happen if Cuda is not installed, we want to raise it after
|
||||
# for any exception other than not having ortmodule, we want to continue
|
||||
# device version validation and raise the exception after.
|
||||
try:
|
||||
from onnxruntime.training.ortmodule._fallback import ORTModuleInitException # noqa: PLC0415
|
||||
|
||||
if isinstance(e, ORTModuleInitException):
|
||||
# ORTModule is present but not ready to run yet
|
||||
has_ortmodule = True
|
||||
except Exception:
|
||||
# ORTModule not present
|
||||
has_ortmodule = False
|
||||
|
||||
if not has_ortmodule:
|
||||
import_ortmodule_exception = e
|
||||
|
||||
# collect onnxruntime package name, version, and cuda version
|
||||
package_name, version, cuda_version = get_package_name_and_version_info()
|
||||
|
||||
if has_ortmodule and cuda_version:
|
||||
try:
|
||||
# collect cuda library build info. the library info may not be available
|
||||
# when the build environment has none or multiple libraries installed
|
||||
try:
|
||||
from .build_and_package_info import cudart_version # noqa: PLC0415
|
||||
except ImportError:
|
||||
warnings.warn("WARNING: failed to get cudart_version from onnxruntime build info.")
|
||||
cudart_version = None
|
||||
|
||||
def print_build_package_info():
|
||||
warnings.warn(f"onnxruntime training package info: package_name: {package_name}")
|
||||
warnings.warn(f"onnxruntime training package info: __version__: {version}")
|
||||
warnings.warn(f"onnxruntime training package info: cuda_version: {cuda_version}")
|
||||
warnings.warn(f"onnxruntime build info: cudart_version: {cudart_version}")
|
||||
|
||||
# collection cuda library info from current environment.
|
||||
from onnxruntime.capi.onnxruntime_collect_build_info import find_cudart_versions # noqa: PLC0415
|
||||
|
||||
local_cudart_versions = find_cudart_versions(build_env=False, build_cuda_version=cuda_version)
|
||||
if cudart_version and local_cudart_versions and cudart_version not in local_cudart_versions:
|
||||
print_build_package_info()
|
||||
warnings.warn("WARNING: failed to find cudart version that matches onnxruntime build info")
|
||||
warnings.warn(f"WARNING: found cudart versions: {local_cudart_versions}")
|
||||
except Exception as e:
|
||||
warnings.warn("WARNING: failed to collect onnxruntime version and build info")
|
||||
print(e)
|
||||
|
||||
if import_ortmodule_exception:
|
||||
raise import_ortmodule_exception
|
||||
|
||||
return has_ortmodule, package_name, version, cuda_version
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
"""
|
||||
Short examples used in the documentation.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
|
||||
def get_example(name):
|
||||
"""
|
||||
Retrieves the absolute file name of an example.
|
||||
"""
|
||||
this = os.path.abspath(os.path.dirname(__file__))
|
||||
full = os.path.join(this, name)
|
||||
if not os.path.exists(full):
|
||||
raise FileNotFoundError(f"Unable to find example '{name}'")
|
||||
return full
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,13 @@
|
||||
backend-test:Q
|
||||
|
||||
xy"Sigmoidtest_sigmoidZ
|
||||
x
|
||||
|
||||
|
||||
|
||||
b
|
||||
y
|
||||
|
||||
|
||||
|
||||
B
|
||||
@@ -0,0 +1,78 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: CalTableFlatBuffers
|
||||
|
||||
import flatbuffers
|
||||
from flatbuffers.compat import import_numpy
|
||||
|
||||
np = import_numpy()
|
||||
|
||||
|
||||
class KeyValue:
|
||||
__slots__ = ["_tab"]
|
||||
|
||||
@classmethod
|
||||
def GetRootAs(cls, buf, offset=0): # noqa: N802
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = KeyValue()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
@classmethod
|
||||
def GetRootAsKeyValue(cls, buf, offset=0): # noqa: N802
|
||||
"""This method is deprecated. Please switch to GetRootAs."""
|
||||
return cls.GetRootAs(buf, offset)
|
||||
|
||||
# KeyValue
|
||||
def Init(self, buf, pos): # noqa: N802
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# KeyValue
|
||||
def Key(self): # noqa: N802
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.String(o + self._tab.Pos)
|
||||
return None
|
||||
|
||||
# KeyValue
|
||||
def Value(self): # noqa: N802
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(6))
|
||||
if o != 0:
|
||||
return self._tab.String(o + self._tab.Pos)
|
||||
return None
|
||||
|
||||
|
||||
def Start(builder): # noqa: N802
|
||||
builder.StartObject(2)
|
||||
|
||||
|
||||
def KeyValueStart(builder): # noqa: N802
|
||||
"""This method is deprecated. Please switch to Start."""
|
||||
return Start(builder)
|
||||
|
||||
|
||||
def AddKey(builder, key): # noqa: N802
|
||||
builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(key), 0)
|
||||
|
||||
|
||||
def KeyValueAddKey(builder, key): # noqa: N802
|
||||
"""This method is deprecated. Please switch to AddKey."""
|
||||
return AddKey(builder, key)
|
||||
|
||||
|
||||
def AddValue(builder, value): # noqa: N802
|
||||
builder.PrependUOffsetTRelativeSlot(1, flatbuffers.number_types.UOffsetTFlags.py_type(value), 0)
|
||||
|
||||
|
||||
def KeyValueAddValue(builder, value): # noqa: N802
|
||||
"""This method is deprecated. Please switch to AddValue."""
|
||||
return AddValue(builder, value)
|
||||
|
||||
|
||||
def End(builder): # noqa: N802
|
||||
return builder.EndObject()
|
||||
|
||||
|
||||
def KeyValueEnd(builder): # noqa: N802
|
||||
"""This method is deprecated. Please switch to End."""
|
||||
return End(builder)
|
||||
@@ -0,0 +1,90 @@
|
||||
# automatically generated by the FlatBuffers compiler, do not modify
|
||||
|
||||
# namespace: CalTableFlatBuffers
|
||||
|
||||
import flatbuffers
|
||||
from flatbuffers.compat import import_numpy
|
||||
|
||||
np = import_numpy()
|
||||
|
||||
|
||||
class TrtTable:
|
||||
__slots__ = ["_tab"]
|
||||
|
||||
@classmethod
|
||||
def GetRootAs(cls, buf, offset=0): # noqa: N802
|
||||
n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
|
||||
x = TrtTable()
|
||||
x.Init(buf, n + offset)
|
||||
return x
|
||||
|
||||
@classmethod
|
||||
def GetRootAsTrtTable(cls, buf, offset=0): # noqa: N802
|
||||
"""This method is deprecated. Please switch to GetRootAs."""
|
||||
return cls.GetRootAs(buf, offset)
|
||||
|
||||
# TrtTable
|
||||
def Init(self, buf, pos): # noqa: N802
|
||||
self._tab = flatbuffers.table.Table(buf, pos)
|
||||
|
||||
# TrtTable
|
||||
def Dict(self, j): # noqa: N802
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
x = self._tab.Vector(o)
|
||||
x += flatbuffers.number_types.UOffsetTFlags.py_type(j) * 4
|
||||
x = self._tab.Indirect(x)
|
||||
from onnxruntime.quantization.CalTableFlatBuffers.KeyValue import KeyValue # noqa: PLC0415
|
||||
|
||||
obj = KeyValue()
|
||||
obj.Init(self._tab.Bytes, x)
|
||||
return obj
|
||||
return None
|
||||
|
||||
# TrtTable
|
||||
def DictLength(self): # noqa: N802
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
if o != 0:
|
||||
return self._tab.VectorLen(o)
|
||||
return 0
|
||||
|
||||
# TrtTable
|
||||
def DictIsNone(self): # noqa: N802
|
||||
o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
|
||||
return o == 0
|
||||
|
||||
|
||||
def Start(builder): # noqa: N802
|
||||
builder.StartObject(1)
|
||||
|
||||
|
||||
def TrtTableStart(builder): # noqa: N802
|
||||
"""This method is deprecated. Please switch to Start."""
|
||||
return Start(builder)
|
||||
|
||||
|
||||
def AddDict(builder, dict): # noqa: N802
|
||||
builder.PrependUOffsetTRelativeSlot(0, flatbuffers.number_types.UOffsetTFlags.py_type(dict), 0)
|
||||
|
||||
|
||||
def TrtTableAddDict(builder, dict): # noqa: N802
|
||||
"""This method is deprecated. Please switch to AddDict."""
|
||||
return AddDict(builder, dict)
|
||||
|
||||
|
||||
def StartDictVector(builder, numElems): # noqa: N802
|
||||
return builder.StartVector(4, numElems, 4)
|
||||
|
||||
|
||||
def TrtTableStartDictVector(builder, numElems): # noqa: N802
|
||||
"""This method is deprecated. Please switch to Start."""
|
||||
return StartDictVector(builder, numElems)
|
||||
|
||||
|
||||
def End(builder): # noqa: N802
|
||||
return builder.EndObject()
|
||||
|
||||
|
||||
def TrtTableEnd(builder): # noqa: N802
|
||||
"""This method is deprecated. Please switch to End."""
|
||||
return End(builder)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,19 @@
|
||||
from .calibrate import ( # noqa: F401
|
||||
CalibraterBase,
|
||||
CalibrationDataReader,
|
||||
CalibrationMethod,
|
||||
MinMaxCalibrater,
|
||||
create_calibrator,
|
||||
)
|
||||
from .qdq_quantizer import QDQQuantizer # noqa: F401
|
||||
from .quant_utils import QuantFormat, QuantType, write_calibration_table # noqa: F401
|
||||
from .quantize import (
|
||||
DynamicQuantConfig, # noqa: F401
|
||||
QuantizationMode, # noqa: F401
|
||||
StaticQuantConfig, # noqa: F401
|
||||
get_qdq_config, # noqa: F401
|
||||
quantize, # noqa: F401
|
||||
quantize_dynamic, # noqa: F401
|
||||
quantize_static, # noqa: F401
|
||||
)
|
||||
from .shape_inference import quant_pre_process # noqa: F401
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,529 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
import onnx.numpy_helper
|
||||
|
||||
try:
|
||||
from onnx.reference.op_run import to_array_extended
|
||||
except ImportError:
|
||||
# old version of onnx.
|
||||
to_array_extended = None
|
||||
|
||||
from .calibrate import TensorData
|
||||
from .onnx_model import ONNXModel
|
||||
from .quant_utils import (
|
||||
DEQUANT_OP_NAME,
|
||||
ONNX_TYPE_TO_NP_TYPE,
|
||||
QUANT_OP_NAME,
|
||||
TENSOR_NAME_QUANT_SUFFIX,
|
||||
find_by_name,
|
||||
get_opset_version,
|
||||
model_has_infer_metadata,
|
||||
normalize_axis,
|
||||
pack_bytes_to_4bit,
|
||||
quantize_data,
|
||||
quantize_nparray,
|
||||
save_and_reload_model_with_shape_infer,
|
||||
tensor_proto_to_array,
|
||||
)
|
||||
from .tensor_quant_overrides import TensorQuantOverridesHelper
|
||||
|
||||
|
||||
class QuantizationParams:
|
||||
def __init__(self, **data: dict[str, Any]):
|
||||
self.data = {}
|
||||
for k, v in data.items():
|
||||
if not isinstance(k, str):
|
||||
raise TypeError(f"Keys must be strings not {type(k)} for k={k!r}.")
|
||||
if k != "axis" and not isinstance(v, (int, str, np.ndarray, float)):
|
||||
raise TypeError(f"Values must be numpy arrays, int, float, str not {type(v)} for k={k!r}.")
|
||||
if k == "axis" and not isinstance(v, int) and v is not None:
|
||||
raise TypeError(f"Axis value must be an int or None, not {type(v)}.")
|
||||
if k == "scale" and v.dtype not in (np.float32, np.float16):
|
||||
raise ValueError(f"scale must a float32 or float16 numpy element but is {v.dtype} for k={k!r}")
|
||||
self.data[k] = v
|
||||
|
||||
def get(self, key, default_value=None):
|
||||
return self.data.get(key, default_value)
|
||||
|
||||
def __iter__(self):
|
||||
yield from self.data
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.data[key]
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self.data[key] = value
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
|
||||
class BaseQuantizer:
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
per_channel,
|
||||
reduce_range,
|
||||
weight_qType,
|
||||
activation_qType,
|
||||
tensors_range,
|
||||
nodes_to_quantize,
|
||||
nodes_to_exclude,
|
||||
op_types_to_quantize,
|
||||
extra_options=None,
|
||||
):
|
||||
if not model_has_infer_metadata(model):
|
||||
model = save_and_reload_model_with_shape_infer(model)
|
||||
self.value_infos = {vi.name: vi for vi in model.graph.value_info}
|
||||
self.value_infos.update({ot.name: ot for ot in model.graph.output})
|
||||
self.value_infos.update({it.name: it for it in model.graph.input})
|
||||
|
||||
self.model = ONNXModel(model)
|
||||
self.opset_version = get_opset_version(model)
|
||||
self.per_channel = per_channel # weight-pack per channel
|
||||
self.reduce_range = reduce_range
|
||||
|
||||
self.extra_options = extra_options if extra_options else {}
|
||||
self.enable_subgraph_quantization = (
|
||||
"EnableSubgraph" in self.extra_options and self.extra_options["EnableSubgraph"]
|
||||
)
|
||||
self.parent = None
|
||||
self.force_quantize_no_input_check = (
|
||||
"ForceQuantizeNoInputCheck" in self.extra_options and self.extra_options["ForceQuantizeNoInputCheck"]
|
||||
)
|
||||
|
||||
# If user does not explicitly set "WeightSymmetric", then the weight's quantization type determines
|
||||
# the symmetry (i.e., signed integer types will use symmetric quantization). See `def is_weight_symmetric()`
|
||||
self._is_weight_symmetric: bool | None = self.extra_options.get("WeightSymmetric", None)
|
||||
self.is_activation_symmetric = self.extra_options.get("ActivationSymmetric", False)
|
||||
self.min_real_range = self.extra_options.get("MinimumRealRange")
|
||||
|
||||
self.activation_qType = getattr(activation_qType, "tensor_type", activation_qType)
|
||||
self.weight_qType = getattr(weight_qType, "tensor_type", weight_qType)
|
||||
|
||||
"""
|
||||
Dictionary specifying the min and max values for tensors. It has following format:
|
||||
{
|
||||
"param_name": [min, max]
|
||||
}
|
||||
example:
|
||||
{
|
||||
'Conv_3:0': [np.float32(0), np.float32(0.5)],
|
||||
'Conv_4:0': [np.float32(1), np.float32(3.5)]
|
||||
}
|
||||
"""
|
||||
if tensors_range is not None and any(not isinstance(t, TensorData) for t in tensors_range.values()):
|
||||
raise TypeError(
|
||||
f"tensors_range contains unexpected types { {type(v) for v in tensors_range.values()} }, not TensorData."
|
||||
)
|
||||
self.tensors_range = tensors_range
|
||||
self.nodes_to_quantize = nodes_to_quantize # specific nodes to quantize
|
||||
self.nodes_to_exclude = nodes_to_exclude # specific nodes to exclude
|
||||
self.op_types_to_quantize = op_types_to_quantize
|
||||
|
||||
# Get tensor-level quantization overrides and ensure they are valid.
|
||||
self.tensor_quant_overrides = TensorQuantOverridesHelper(self.extra_options.get("TensorQuantOverrides", {}))
|
||||
|
||||
self.initializers = {initzer.name: initzer for initzer in self.model.initializer()}
|
||||
overrides_valid, overrides_err = self.tensor_quant_overrides.is_valid(
|
||||
self.initializers, self.value_infos.keys(), activation_qType
|
||||
)
|
||||
if not overrides_valid:
|
||||
raise ValueError(overrides_err)
|
||||
|
||||
self.tensor_quant_override_qtypes = self.tensor_quant_overrides.get_quant_types()
|
||||
|
||||
def is_weight_symmetric(self, weight_quant_type: onnx.TensorProto.DataType) -> bool:
|
||||
if self._is_weight_symmetric is not None:
|
||||
return self._is_weight_symmetric # Return value explicitly set by user.
|
||||
return weight_quant_type in (
|
||||
onnx.TensorProto.INT4,
|
||||
onnx.TensorProto.INT8,
|
||||
onnx.TensorProto.INT16,
|
||||
onnx.TensorProto.FLOAT8E4M3FN,
|
||||
)
|
||||
|
||||
def quantize_model(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def is_input_a_initializer(self, input_name):
|
||||
initializer = find_by_name(input_name, self.model.initializer())
|
||||
return initializer is not None
|
||||
|
||||
def is_per_channel(self):
|
||||
return self.per_channel
|
||||
|
||||
def is_valid_quantize_weight(self, weight_name):
|
||||
weight = find_by_name(weight_name, self.model.initializer())
|
||||
if weight is not None:
|
||||
return weight.data_type in (onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT16)
|
||||
if (not self.enable_subgraph_quantization) or (self.parent is None):
|
||||
return False
|
||||
return self.parent.is_valid_quantize_weight(weight_name)
|
||||
|
||||
def should_quantize_node(self, node):
|
||||
if (
|
||||
self.nodes_to_quantize is not None
|
||||
and len(self.nodes_to_quantize) != 0
|
||||
and node.name not in self.nodes_to_quantize
|
||||
):
|
||||
return False
|
||||
|
||||
if node.op_type not in self.op_types_to_quantize:
|
||||
return False
|
||||
|
||||
if node.op_type in (DEQUANT_OP_NAME, QUANT_OP_NAME):
|
||||
return False
|
||||
|
||||
if self.nodes_to_exclude is not None and node.name in self.nodes_to_exclude:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def quantize_bias_static_impl(self, bias_name, input_scale, weight_scale, beta=1.0):
|
||||
"""
|
||||
Quantized the bias. Zero Point == 0 and Scale == Input_Scale * Weight_Scale
|
||||
"""
|
||||
|
||||
# get bias
|
||||
bias_initializer = find_by_name(bias_name, self.model.initializer())
|
||||
bias_data = tensor_proto_to_array(bias_initializer)
|
||||
quantized_bias_name = bias_name + TENSOR_NAME_QUANT_SUFFIX
|
||||
|
||||
# quantize bias
|
||||
if self.weight_qType == onnx.TensorProto.FLOAT8E4M3FN:
|
||||
data = np.asarray(bias_data)
|
||||
if data.dtype == np.float16:
|
||||
node_qtype = onnx.TensorProto.FLOAT16
|
||||
elif data.dtype == np.float32:
|
||||
node_qtype = onnx.TensorProto.FLOAT
|
||||
else:
|
||||
raise TypeError(f"Only float16 or float32 are supported with float 8 but bias dtype is {data.dtype}.")
|
||||
quantized_data = data.astype(np.float32)
|
||||
bias_scale = np.array([1], dtype=quantized_data.dtype)
|
||||
bias_scale_data = bias_scale.reshape(-1)
|
||||
packed_bias_initializer = onnx.numpy_helper.from_array(quantized_data, quantized_bias_name)
|
||||
self.model.initializer_extend([packed_bias_initializer])
|
||||
node_type = "Cast"
|
||||
else:
|
||||
# calculate scale for bias
|
||||
# TODO: This formula should be explained including why the scale is not estimated for the bias as well.
|
||||
bias_scale = input_scale * weight_scale * beta
|
||||
|
||||
# Quantize by dividing by bias_scale
|
||||
quantized_data = np.asarray(bias_data, dtype=np.float64) / np.asarray(bias_scale, dtype=np.float64)
|
||||
quantized_data = quantized_data.round()
|
||||
|
||||
# Clip quantized data to the range of a int32
|
||||
int32_min = np.float64(np.iinfo(np.int32).min)
|
||||
int32_max = np.float64(np.iinfo(np.int32).max)
|
||||
if np.any(quantized_data < int32_min) or np.any(quantized_data > int32_max):
|
||||
logging.warning(
|
||||
f"Quantized bias `{bias_name}` exceeds the range of a int32. The bias scale is too small."
|
||||
)
|
||||
|
||||
quantized_data = np.clip(quantized_data, int32_min, int32_max).astype(np.int32)
|
||||
|
||||
# update bias initializer
|
||||
bias_np_data = np.asarray(quantized_data, dtype=np.int32).reshape(bias_initializer.dims)
|
||||
packed_bias_initializer = onnx.numpy_helper.from_array(bias_np_data, quantized_bias_name)
|
||||
self.model.initializer_extend([packed_bias_initializer])
|
||||
|
||||
# Bias's scale dtype should match the original bias data's unquantized type (float32 or float16).
|
||||
bias_scale_data = np.asarray(bias_scale, dtype=bias_data.dtype).reshape(-1)
|
||||
node_type = "DequantizeLinear"
|
||||
node_qtype = self.weight_qType
|
||||
|
||||
# update scale initializer
|
||||
quantized_bias_scale_name = quantized_bias_name + "_scale"
|
||||
packed_bias_scale_initializer = onnx.numpy_helper.from_array(bias_scale_data, quantized_bias_scale_name)
|
||||
self.model.initializer_extend([packed_bias_scale_initializer])
|
||||
|
||||
# update zero initializer
|
||||
if self.weight_qType == onnx.TensorProto.FLOAT8E4M3FN:
|
||||
tensor_type = self.weight_qType
|
||||
else:
|
||||
tensor_type = onnx.TensorProto.INT32
|
||||
|
||||
quantized_bias_zp_name = quantized_bias_name + "_zero_point"
|
||||
if self.weight_qType == onnx.TensorProto.FLOAT8E4M3FN:
|
||||
packed_bias_zp_initializer = onnx.helper.make_tensor(quantized_bias_zp_name, self.weight_qType, [1], [0.0])
|
||||
elif bias_scale.size > 1:
|
||||
bias_zp_data = np.zeros(bias_scale.shape, dtype=np.int32).reshape(-1)
|
||||
packed_bias_zp_initializer = onnx.numpy_helper.from_array(bias_zp_data, quantized_bias_zp_name)
|
||||
else:
|
||||
packed_bias_zp_initializer = onnx.helper.make_tensor(quantized_bias_zp_name, tensor_type, [], [0])
|
||||
self.model.initializer_extend([packed_bias_zp_initializer])
|
||||
|
||||
return (
|
||||
quantized_bias_name,
|
||||
quantized_bias_scale_name,
|
||||
quantized_bias_zp_name,
|
||||
bias_scale_data,
|
||||
node_type,
|
||||
node_qtype,
|
||||
)
|
||||
|
||||
def quantize_initializer_impl(self, weight, qType, reduce_range=False, keep_float_weight=False):
|
||||
"""
|
||||
:param weight: TensorProto initializer
|
||||
:param qType: type to quantize to
|
||||
:param keep_float_weight: Whether to quantize the weight. In some cases, we only want to qunatize scale and zero point.
|
||||
If keep_float_weight is False, quantize the weight, or don't quantize the weight.
|
||||
:return: quantized weight name, zero point name, scale name
|
||||
"""
|
||||
# TODO(adrianlizarraga): This function is now only used by onnx_quantizer.py, so move it there.
|
||||
q_weight_name = weight.name + TENSOR_NAME_QUANT_SUFFIX
|
||||
zp_name = weight.name + "_zero_point"
|
||||
scale_name = weight.name + "_scale"
|
||||
|
||||
# Quantize weight data. Use quantization overrides if provided by the user.
|
||||
weight_data = tensor_proto_to_array(weight)
|
||||
quant_overrides = self.tensor_quant_overrides.get_per_tensor_overrides(weight.name, default_val={})
|
||||
if "quant_type" in quant_overrides:
|
||||
qType = quant_overrides["quant_type"].tensor_type # noqa: N806
|
||||
|
||||
if "scale" in quant_overrides and "zero_point" in quant_overrides:
|
||||
zero_point = np.array(quant_overrides["zero_point"], dtype=ONNX_TYPE_TO_NP_TYPE[qType])
|
||||
scale = np.array(quant_overrides["scale"])
|
||||
q_weight_data = quantize_nparray(qType, weight_data.flatten(), scale, zero_point)
|
||||
assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}"
|
||||
assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, (
|
||||
f"Unexpected dtype {zero_point.dtype}"
|
||||
)
|
||||
assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}"
|
||||
|
||||
else:
|
||||
symmetric = self.is_weight_symmetric(qType) if qType == self.weight_qType else self.is_activation_symmetric
|
||||
zero_point, scale, q_weight_data = quantize_data(
|
||||
weight_data.flatten(),
|
||||
qType,
|
||||
quant_overrides.get("symmetric", symmetric),
|
||||
reduce_range=quant_overrides.get("reduce_range", self.reduce_range and reduce_range),
|
||||
min_real_range=self.min_real_range,
|
||||
rmin_override=quant_overrides.get("rmin"),
|
||||
rmax_override=quant_overrides.get("rmax"),
|
||||
)
|
||||
|
||||
assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}"
|
||||
assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, (
|
||||
f"Unexpected dtype {zero_point.dtype}"
|
||||
)
|
||||
assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}"
|
||||
|
||||
scale_dtype = weight.data_type
|
||||
scale_initializer = onnx.helper.make_tensor(scale_name, scale_dtype, [], scale.reshape((-1,)).tolist())
|
||||
zero_initializer = onnx.helper.make_tensor(zp_name, qType, [], zero_point.reshape((-1,)).tolist())
|
||||
self.model.initializer_extend([scale_initializer, zero_initializer])
|
||||
|
||||
if not keep_float_weight:
|
||||
if self.weight_qType == onnx.TensorProto.FLOAT8E4M3FN:
|
||||
q_weight_initializer = onnx.TensorProto()
|
||||
q_weight_initializer.data_type = self.weight_qType
|
||||
q_weight_initializer.dims.extend(weight.dims)
|
||||
q_weight_initializer.name = q_weight_name
|
||||
# Do not remove .flatten().copy() numpy is not clear about data persistence.
|
||||
q_weight_initializer.raw_data = q_weight_data.flatten().copy().tobytes()
|
||||
if to_array_extended is not None:
|
||||
# This test should not be needed but it helped catch some issues
|
||||
# with data persistence and tobytes.
|
||||
check = to_array_extended(q_weight_initializer)
|
||||
if check.shape != weight_data.shape or check.tobytes() != q_weight_data.tobytes():
|
||||
raise RuntimeError(
|
||||
f"The initializer of shape {weight_data.shape} could not be created, expecting "
|
||||
f"{q_weight_data.tobytes()[:10]}, got {check.tobytes()[:10]} and shape={weight.shape}"
|
||||
f"\nraw={str(q_weight_initializer)[:200]}."
|
||||
)
|
||||
elif qType in (onnx.TensorProto.INT4, onnx.TensorProto.UINT4):
|
||||
if q_weight_data.dtype not in (np.int8, np.uint8):
|
||||
raise RuntimeError(
|
||||
f"Quantized weights for {q_weight_name} must be 8-bit before packing as 4-bit values."
|
||||
)
|
||||
|
||||
# We do not use onnx.helper.pack_float32_to_4bit() due to performance.
|
||||
# This can be the difference between a large model taking 30 minutes to quantize vs 5 minutes.
|
||||
packed_data = bytes(pack_bytes_to_4bit(q_weight_data.tobytes()))
|
||||
|
||||
# We only use onnx.helper.make_tensor with raw data due to bug: https://github.com/onnx/onnx/pull/6161
|
||||
q_weight_initializer = onnx.helper.make_tensor(q_weight_name, qType, weight.dims, packed_data, raw=True)
|
||||
else:
|
||||
q_weight_data = np.asarray(q_weight_data, dtype=onnx.helper.tensor_dtype_to_np_dtype(qType)).reshape(
|
||||
weight.dims
|
||||
)
|
||||
q_weight_initializer = onnx.numpy_helper.from_array(q_weight_data, q_weight_name)
|
||||
self.model.initializer_extend([q_weight_initializer])
|
||||
|
||||
return q_weight_name, zp_name, scale_name
|
||||
|
||||
def quantize_weight_per_channel_impl(
|
||||
self,
|
||||
weight_name,
|
||||
weight_qType,
|
||||
channel_axis,
|
||||
reduce_range=True,
|
||||
keep_float_weight=False,
|
||||
):
|
||||
# TODO(adrianlizarraga): This function is now only used by onnx_quantizer.py, so move it there.
|
||||
initializer = find_by_name(weight_name, self.model.initializer())
|
||||
if initializer is None:
|
||||
raise ValueError("{} is not an initializer", weight_name)
|
||||
|
||||
weights = tensor_proto_to_array(initializer)
|
||||
weights_rank = len(weights.shape)
|
||||
is_axis_valid, axis_norm = normalize_axis(channel_axis, weights_rank)
|
||||
if not is_axis_valid:
|
||||
raise ValueError(
|
||||
f"Weight {weight_name} has a per-channel axis with value {channel_axis} that is "
|
||||
f"out-of-bounds for rank {weights_rank}"
|
||||
)
|
||||
|
||||
channel_axis = axis_norm
|
||||
channel_count = weights.shape[channel_axis]
|
||||
quant_overrides_for_channels = self.tensor_quant_overrides.get_per_channel_overrides(
|
||||
weight_name, default_val=[{"axis": channel_axis}]
|
||||
)
|
||||
|
||||
num_channel_overrides = len(quant_overrides_for_channels)
|
||||
if num_channel_overrides != 1 and num_channel_overrides != channel_count:
|
||||
raise ValueError(
|
||||
f"Per-channel tensor quantization overrides for {weight_name} must have "
|
||||
f"either 1 or {channel_count} elements in the list of dictionaries."
|
||||
)
|
||||
|
||||
is_axis_override_valid, axis_override = normalize_axis(quant_overrides_for_channels[0]["axis"], weights_rank)
|
||||
if not is_axis_override_valid or axis_override != channel_axis:
|
||||
raise ValueError(
|
||||
f"Tensor quantization overrides for {weight_name} specify an unexpected axis. "
|
||||
f"Expected {channel_axis}, but got {quant_overrides_for_channels[0]['axis']}."
|
||||
)
|
||||
|
||||
# If user provides per-channel quantization overrides, all channels must use the same quant_type,
|
||||
# axis, symmetric, and reduce_range values. So, just use the first channel's values.
|
||||
if "quant_type" in quant_overrides_for_channels[0]:
|
||||
weight_qType = quant_overrides_for_channels[0]["quant_type"].tensor_type # noqa: N806
|
||||
|
||||
symmetric = quant_overrides_for_channels[0].get("symmetric", self.is_weight_symmetric(weight_qType))
|
||||
reduce_range = quant_overrides_for_channels[0].get("reduce_range", self.reduce_range and reduce_range)
|
||||
zero_point_list = []
|
||||
scale_list = []
|
||||
quantized_per_channel_data_list = []
|
||||
weights_shape = list(weights.shape)
|
||||
reshape_dims = list(weights_shape) # deep copy
|
||||
reshape_dims[channel_axis] = 1 # only one per channel for reshape
|
||||
for i in range(channel_count):
|
||||
per_channel_data = weights.take(i, channel_axis)
|
||||
channel_override_index = i if i < num_channel_overrides else 0
|
||||
channel_quant_overrides = quant_overrides_for_channels[channel_override_index]
|
||||
|
||||
if "scale" in channel_quant_overrides and "zero_point" in channel_quant_overrides:
|
||||
zero_point = np.array(channel_quant_overrides["zero_point"], dtype=ONNX_TYPE_TO_NP_TYPE[weight_qType])
|
||||
scale = np.array(channel_quant_overrides["scale"])
|
||||
quantized_per_channel_data = quantize_nparray(
|
||||
weight_qType, per_channel_data.flatten(), scale, zero_point
|
||||
)
|
||||
assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}"
|
||||
assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, (
|
||||
f"Unexpected dtype {zero_point.dtype}"
|
||||
)
|
||||
assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}"
|
||||
assert isinstance(quantized_per_channel_data, np.ndarray), (
|
||||
f"Unexpected type {type(quantized_per_channel_data)}"
|
||||
)
|
||||
|
||||
else:
|
||||
zero_point, scale, quantized_per_channel_data = quantize_data(
|
||||
per_channel_data.flatten(),
|
||||
weight_qType,
|
||||
symmetric,
|
||||
reduce_range=reduce_range,
|
||||
min_real_range=self.min_real_range,
|
||||
rmin_override=channel_quant_overrides.get("rmin"),
|
||||
rmax_override=channel_quant_overrides.get("rmax"),
|
||||
)
|
||||
|
||||
assert isinstance(zero_point, np.ndarray), f"Unexpected type {type(zero_point)}"
|
||||
assert zero_point.dtype != np.float32 and zero_point.dtype != np.float16, (
|
||||
f"Unexpected dtype {zero_point.dtype}"
|
||||
)
|
||||
assert isinstance(scale, np.ndarray), f"Unexpected type {type(scale)}"
|
||||
assert isinstance(quantized_per_channel_data, np.ndarray), (
|
||||
f"Unexpected type {type(quantized_per_channel_data)}"
|
||||
)
|
||||
|
||||
zero_point_list.append(zero_point)
|
||||
scale_list.append(scale)
|
||||
quantized_per_channel_data_list.append(np.asarray(quantized_per_channel_data).reshape(reshape_dims))
|
||||
|
||||
# combine per_channel_data into one
|
||||
quantized_weights = np.concatenate(quantized_per_channel_data_list, channel_axis)
|
||||
q_weight_name = weight_name + TENSOR_NAME_QUANT_SUFFIX
|
||||
zp_name = weight_name + "_zero_point"
|
||||
scale_name = weight_name + "_scale"
|
||||
|
||||
# Update packed weight, zero point, and scale initializers
|
||||
zero_scale_shape = [initializer.dims[channel_axis]]
|
||||
scale_initializer = onnx.helper.make_tensor(
|
||||
scale_name, initializer.data_type, zero_scale_shape, np.hstack(scale_list).tolist()
|
||||
)
|
||||
zero_initializer = onnx.helper.make_tensor(
|
||||
zp_name, weight_qType, zero_scale_shape, np.hstack(zero_point_list).tolist()
|
||||
)
|
||||
|
||||
self.model.initializer_extend([scale_initializer, zero_initializer])
|
||||
|
||||
if not keep_float_weight:
|
||||
if weight_qType in (onnx.TensorProto.INT4, onnx.TensorProto.UINT4):
|
||||
if quantized_weights.dtype not in (np.int8, np.uint8):
|
||||
raise RuntimeError(
|
||||
f"Quantized weights for {q_weight_name} must be 8-bit before packing as 4-bit values."
|
||||
)
|
||||
|
||||
# We do not use onnx.helper.pack_float32_to_4bit() due to performance.
|
||||
# This can be the difference between a large model taking 30 minutes to quantize vs 5 minutes.
|
||||
packed_data = bytes(pack_bytes_to_4bit(quantized_weights.tobytes()))
|
||||
|
||||
# We only use onnx.helper.make_tensor with raw data due to bug: https://github.com/onnx/onnx/pull/6161
|
||||
q_weight_initializer = onnx.helper.make_tensor(
|
||||
q_weight_name, weight_qType, weights_shape, packed_data, raw=True
|
||||
)
|
||||
self.model.initializer_extend([q_weight_initializer])
|
||||
else:
|
||||
quantized_weights = np.asarray(
|
||||
quantized_weights,
|
||||
dtype=onnx.helper.tensor_dtype_to_np_dtype(weight_qType),
|
||||
).reshape(initializer.dims)
|
||||
q_weight_initializer = onnx.numpy_helper.from_array(quantized_weights, q_weight_name)
|
||||
self.model.initializer_extend([q_weight_initializer])
|
||||
|
||||
return q_weight_name, zp_name, scale_name
|
||||
|
||||
def adjust_tensor_ranges(self):
|
||||
if self.tensors_range is None:
|
||||
return
|
||||
|
||||
for node in self.model.nodes():
|
||||
# adjust tensor_ranges for input of Clip and Relu node
|
||||
if node.op_type in ["Clip", "Relu"]:
|
||||
if not self.should_quantize_node(node):
|
||||
continue
|
||||
if len(self.model.input_name_to_nodes()[node.input[0]]) != 1:
|
||||
continue
|
||||
if node.input[0] not in self.tensors_range or node.output[0] not in self.tensors_range:
|
||||
continue
|
||||
td = self.tensors_range[node.output[0]]
|
||||
if not isinstance(td, TensorData):
|
||||
raise TypeError(f"Unexpected type {type(td)} for {node.output[0]!r}.")
|
||||
self.tensors_range[node.input[0]] = td
|
||||
# Adjust Softmax to range from 0.0 to 1.0
|
||||
elif node.op_type == "Softmax":
|
||||
if not self.should_quantize_node(node):
|
||||
continue
|
||||
self.tensors_range[node.output[0]] = TensorData(lowest=np.float32(0.0), highest=np.float32(1.0))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,2 @@
|
||||
from .preprocess import qnn_preprocess_model # noqa: F401
|
||||
from .quant_config import get_qnn_qdq_config # noqa: F401
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,132 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
import onnx
|
||||
|
||||
from ...fusions import Fusion
|
||||
from ...onnx_model import ONNXModel
|
||||
|
||||
|
||||
class FusionLpNormalization(Fusion):
|
||||
def __init__(self, model: ONNXModel, epsilon: float = 1e-12):
|
||||
super().__init__(model, "LpNormalization", "ReduceL2")
|
||||
self.epsilon = epsilon
|
||||
|
||||
def fuse(
|
||||
self,
|
||||
reduce_node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
):
|
||||
"""
|
||||
Interface function that tries to fuse a node sequence containing a ReduceL2 node into a single
|
||||
LpNormalization node.
|
||||
|
||||
Pattern 1:
|
||||
[root] --> ReduceL2 -----> Clip --> Expand ----> Div -->
|
||||
| (axis=-1) (min=epsilon) (shape=root) ^
|
||||
| (keepdims=True) |
|
||||
| |
|
||||
+-----------------------------------------------+
|
||||
Notes:
|
||||
- ReduceL2 must use the last axis, and keepdims == True
|
||||
- Clip must only have a min attribute that is ~1e-12
|
||||
- Expand must restore the shape to root.shape
|
||||
- The output of Expand must be the second input to Div.
|
||||
"""
|
||||
if reduce_node.output[0] not in input_name_to_nodes:
|
||||
return
|
||||
|
||||
# ReduceL2 must have one Clip child
|
||||
children = input_name_to_nodes[reduce_node.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Clip":
|
||||
return
|
||||
|
||||
# ReduceL2 must have keepdims == True
|
||||
keepdims = self.get_node_attribute(reduce_node, "keepdims")
|
||||
if not keepdims:
|
||||
return
|
||||
|
||||
# ReduceL2 axes must refer only to the last dimension.
|
||||
# Axes became an input in opset 18. Before then, axes was an attribute
|
||||
reduce_input_ttype = self.model.get_tensor_type(reduce_node.input[0])
|
||||
if not reduce_input_ttype:
|
||||
return
|
||||
|
||||
reduce_input_shape = self.tensor_shape_to_list(reduce_input_ttype)
|
||||
if not reduce_input_shape:
|
||||
return
|
||||
|
||||
axes = self.get_node_attribute(reduce_node, "axes")
|
||||
if not axes and len(reduce_node.input) > 1:
|
||||
axes = self.model.get_constant_value(reduce_node.input[1])
|
||||
|
||||
if not axes or len(axes) != 1:
|
||||
return
|
||||
|
||||
last_dim = len(reduce_input_shape) - 1
|
||||
if axes[0] != -1 and axes[0] != last_dim:
|
||||
return
|
||||
|
||||
# Clip node must have a min attribute approximately equal to 1e-12
|
||||
clip_node = children[0]
|
||||
clip_min = self.get_node_attribute(clip_node, "min")
|
||||
if clip_min is None and len(clip_node.input) > 1:
|
||||
clip_min = self.model.get_constant_value(clip_node.input[1])
|
||||
|
||||
clip_max = self.get_node_attribute(clip_node, "max") # TODO: clip_max could be FLOAT_MAX
|
||||
if clip_max is None and len(clip_node.input) > 2:
|
||||
clip_max = self.model.get_constant_value(clip_node.input[2])
|
||||
|
||||
if not (clip_max is None and clip_min is not None and clip_min > 0 and abs(clip_min - self.epsilon) < 1e-13):
|
||||
return
|
||||
|
||||
if clip_node.output[0] not in input_name_to_nodes:
|
||||
return
|
||||
|
||||
# Clip must have a single Expand child.
|
||||
children = input_name_to_nodes[clip_node.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Expand":
|
||||
return
|
||||
|
||||
expand_node = children[0]
|
||||
if expand_node.output[0] not in input_name_to_nodes:
|
||||
return
|
||||
|
||||
# Expand must have a single Div child
|
||||
children = input_name_to_nodes[expand_node.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Div":
|
||||
return
|
||||
|
||||
div_node = children[0]
|
||||
|
||||
# The first input to Div must be the root of the subgraph (i.e., reduce_node.input[0])
|
||||
# The second input to Div must be the output of the Expand.
|
||||
# As long as these two inputs go to the same Div node, then ONNX validation will ensure that
|
||||
# their shapes match.
|
||||
if div_node.input[0] != reduce_node.input[0]:
|
||||
return
|
||||
if div_node.input[1] != expand_node.output[0]:
|
||||
return
|
||||
|
||||
subgraph_input = reduce_node.input[0]
|
||||
subgraph_output = div_node.output[0]
|
||||
|
||||
subgraph_nodes = [reduce_node, clip_node, expand_node, div_node]
|
||||
if not self.is_safe_to_fuse_nodes(subgraph_nodes, [subgraph_output], input_name_to_nodes, output_name_to_node):
|
||||
return
|
||||
|
||||
self.nodes_to_remove.extend(subgraph_nodes)
|
||||
fused_node = onnx.helper.make_node(
|
||||
self.fused_op_type,
|
||||
name=self.create_unique_node_name(),
|
||||
inputs=[subgraph_input],
|
||||
outputs=[subgraph_output],
|
||||
p=2,
|
||||
axis=-1,
|
||||
)
|
||||
self.nodes_to_add.append(fused_node)
|
||||
@@ -0,0 +1,162 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
"""Define SpaceToDepth fusion."""
|
||||
|
||||
import onnx
|
||||
|
||||
from ... import fusions, onnx_model
|
||||
|
||||
|
||||
class FusionSpaceToDepth(fusions.Fusion):
|
||||
"""Fusion for SpaceToDepth."""
|
||||
|
||||
def __init__(self, model: onnx_model.ONNXModel):
|
||||
"""Initialize.
|
||||
|
||||
Args:
|
||||
model: An onnx_model.ONNXModel instance.
|
||||
"""
|
||||
super().__init__(model, "SpaceToDepth", "Reshape")
|
||||
|
||||
def _fuse_yolo(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
):
|
||||
"""Fuse for early version of YOLO.
|
||||
|
||||
Pattern:
|
||||
|
||||
| [N, C, H, W]
|
||||
Reshape
|
||||
| [N, C, H/blk, blk, W/blk, blk]
|
||||
Transpose
|
||||
| [N, C, H/blk, W/blk, blk, blk]
|
||||
Reshape
|
||||
| [N, C, H/blk * W/blk, blk * blk]
|
||||
Transpose
|
||||
| [N, C, blk * blk, H/blk * W/blk]
|
||||
Reshape
|
||||
| [N, C, blk * blk, H/blk, W/blk]
|
||||
Transpose
|
||||
| [N, blk * blk, C, H/blk, W/blk]
|
||||
Reshape
|
||||
| [N, blk * blk * C, H/blk, W/blk]
|
||||
|
||||
This sequence can be fused into a single SpaceToDepth with blocksize `blk`. Note that unlike DepthToSpace
|
||||
supporting DCR or CRD mode, SpaceToDepth only supports DCR mode in its latest opset version (13), which matches
|
||||
the pattern here.
|
||||
"""
|
||||
reshape_node1 = node
|
||||
|
||||
def get_target_child(parent_node, target_op_type):
|
||||
"""Get target child of given node."""
|
||||
if parent_node.output[0] not in input_name_to_nodes:
|
||||
return None
|
||||
|
||||
children = input_name_to_nodes[parent_node.output[0]]
|
||||
if len(children) > 1 or children[0].op_type != target_op_type:
|
||||
return None
|
||||
|
||||
return children[0]
|
||||
|
||||
if (
|
||||
(transpose_node1 := get_target_child(reshape_node1, "Transpose")) is None
|
||||
or (reshape_node2 := get_target_child(transpose_node1, "Reshape")) is None
|
||||
or (transpose_node2 := get_target_child(reshape_node2, "Transpose")) is None
|
||||
or (reshape_node3 := get_target_child(transpose_node2, "Reshape")) is None
|
||||
or (transpose_node3 := get_target_child(reshape_node3, "Transpose")) is None
|
||||
or (reshape_node4 := get_target_child(transpose_node3, "Reshape")) is None
|
||||
):
|
||||
return False
|
||||
|
||||
def get_tensor_shape(tensor_name):
|
||||
"""Get shape for given tensor name."""
|
||||
tensor_type = self.model.get_tensor_type(tensor_name)
|
||||
if not tensor_type:
|
||||
return None
|
||||
|
||||
tensor_shape = self.tensor_shape_to_list(tensor_type)
|
||||
if not tensor_shape:
|
||||
return None
|
||||
|
||||
return tensor_shape
|
||||
|
||||
if (
|
||||
(input_shape := get_tensor_shape(reshape_node1.input[0])) is None
|
||||
or (reshape_shape1 := get_tensor_shape(reshape_node1.output[0])) is None
|
||||
or (reshape_shape2 := get_tensor_shape(reshape_node2.output[0])) is None
|
||||
or (reshape_shape3 := get_tensor_shape(reshape_node3.output[0])) is None
|
||||
or (reshape_shape4 := get_tensor_shape(reshape_node4.output[0])) is None
|
||||
):
|
||||
return False
|
||||
|
||||
transpose_perm1 = self.get_node_attribute(transpose_node1, "perm")
|
||||
transpose_perm2 = self.get_node_attribute(transpose_node2, "perm")
|
||||
transpose_perm3 = self.get_node_attribute(transpose_node3, "perm")
|
||||
|
||||
# Check rank.
|
||||
if (
|
||||
len(input_shape) != 4
|
||||
or len(reshape_shape1) != 6
|
||||
or len(reshape_shape2) != 4
|
||||
or len(reshape_shape3) != 5
|
||||
or len(reshape_shape4) != 4
|
||||
):
|
||||
return False
|
||||
|
||||
# Check shape and perm.
|
||||
batch, channel, height, width = input_shape
|
||||
blocksize = reshape_shape1[3]
|
||||
if (
|
||||
reshape_shape1 != [batch, channel, height // blocksize, blocksize, width // blocksize, blocksize]
|
||||
or transpose_perm1 != [0, 1, 2, 4, 3, 5]
|
||||
or reshape_shape2 != [batch, channel, (height // blocksize) * (width // blocksize), blocksize**2]
|
||||
or transpose_perm2 != [0, 1, 3, 2]
|
||||
or reshape_shape3 != [batch, channel, blocksize**2, height // blocksize, width // blocksize]
|
||||
or transpose_perm3 != [0, 2, 1, 3, 4]
|
||||
or reshape_shape4 != [batch, blocksize**2 * channel, height // blocksize, width // blocksize]
|
||||
):
|
||||
return False
|
||||
|
||||
self.nodes_to_remove.extend(
|
||||
[
|
||||
reshape_node1,
|
||||
transpose_node1,
|
||||
reshape_node2,
|
||||
transpose_node2,
|
||||
reshape_node3,
|
||||
transpose_node3,
|
||||
reshape_node4,
|
||||
]
|
||||
)
|
||||
|
||||
s2d_node = onnx.helper.make_node(
|
||||
self.fused_op_type,
|
||||
name=self.create_unique_node_name(),
|
||||
inputs=[reshape_node1.input[0]],
|
||||
outputs=[reshape_node4.output[0]],
|
||||
blocksize=blocksize,
|
||||
)
|
||||
self.nodes_to_add.append(s2d_node)
|
||||
|
||||
return True
|
||||
|
||||
def fuse(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
):
|
||||
"""Fuse a sequence of Reshape and Transpose nodes into a single SpaceToDepth node.
|
||||
|
||||
Args:
|
||||
node: An onnx.NodeProto matching the specified search type (i.e., Reshape).
|
||||
input_name_to_nodes: A dict mapping tensor name to consumed nodes.
|
||||
output_name_to_node: A dict mapping tensor name to produced node.
|
||||
"""
|
||||
self._fuse_yolo(node, input_name_to_nodes, output_name_to_node)
|
||||
@@ -0,0 +1,413 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
import onnx
|
||||
|
||||
from ...quant_utils import QuantType
|
||||
from ...tensor_quant_overrides import QuantTypeInfo, TensorQuantOverridesHelper
|
||||
|
||||
|
||||
@dataclass
|
||||
class TensorTypeRequest:
|
||||
"""
|
||||
Bundles desired quantization type requests for a tensor. A distinction is made between the
|
||||
produced type and the consumed type.
|
||||
"""
|
||||
|
||||
# The tensor's quant type at the producer end. If None, assumed to be the default activation quant type.
|
||||
producer: QuantTypeInfo | None
|
||||
|
||||
# The tensor's quant type received by a set of consumer nodes.
|
||||
# If None, assumed to be the default activation quant type for all consumers.
|
||||
# consumers[1] is a set of consumer node names.
|
||||
consumers: tuple[QuantTypeInfo, set[str]] | None
|
||||
|
||||
|
||||
class MixedPrecisionTensorQuantOverridesFixer:
|
||||
"""
|
||||
Helper that generates tensor quantization overrides for mixed-precision QDQ models.
|
||||
|
||||
Specifically, this helper fixes an initial set of quantization overrides that assign a non-default
|
||||
activation quantization type to one or more tensors by doing the following:
|
||||
- Inferring which other tensors need to be overridden to the non-default activation quantization type.
|
||||
- Inserting quantization data type conversions.
|
||||
|
||||
Example:
|
||||
--------
|
||||
|
||||
Float model:
|
||||
|
||||
input_0 --> Op1 --> Op3 --> Op5 --> Op6 --> output_0
|
||||
^
|
||||
|
|
||||
input_1 --> Op2 -+-> Op4 ----+
|
||||
|
|
||||
+-> Op7 --> output_1
|
||||
|
|
||||
+-> Op8 --> output_2
|
||||
|
||||
If we'd like to quantize this model to uint8 precision, but would like to make sure tensor "Op4_out"
|
||||
is quantized to 16-bit, then we would specify the following initial tensor quantization overrides:
|
||||
|
||||
```
|
||||
init_overrides = {"Op4_out": [{"quant_type": QuantType.QUInt16}]}
|
||||
```
|
||||
|
||||
These initial overrides may not create a valid model because Op4 and Op5 may require both the input and output
|
||||
to be the same type (e.g., uint16). This helper fixes the overrides so that input/output data types
|
||||
are valid:
|
||||
|
||||
```
|
||||
overrides = TensorQuantOverridesHelper(init_overrides)
|
||||
|
||||
fixer = MixedPrecisionTensorQuantOverridesFixer.create_from_model(overrides, model, QuantType.QUInt8)
|
||||
fixer.apply(
|
||||
default_activation_qtype=QuantType.QUInt8,
|
||||
default_activation_symmetric=False,
|
||||
)
|
||||
```
|
||||
|
||||
The above snippet generates the following "fixed" overrides (get via overrides.get_dict()):
|
||||
|
||||
{
|
||||
"Op2_out": [{"quant_type": QUInt8, "convert": {"quant_type": QUInt16, "recv_nodes": {"Op4"}}}],
|
||||
"Op3_out": [{"quant_type": QUInt8, "convert": {"quant_type": QUInt16, "recv_nodes": {"Op5"}}}],
|
||||
"Op4_out": [{"quant_type": QUInt16}],
|
||||
"Op5_out": [{"quant_type": QUInt16, "convert": {"quant_type": QUInt8, "recv_nodes": {"Op6"}}}]
|
||||
}
|
||||
|
||||
How to interpret the fixed overrides:
|
||||
- Op2's output is consumed by Op4, Op7, and Op8. Op4 consumes the converted u16 type,
|
||||
but Op7 and Op8 consume the original u8 type.
|
||||
- Op3's output is converted from u8 to u16. Op5 consumes the converted u16 type.
|
||||
- Op4's output is just u16 (not converted). All consumers of Op4_out get the u16 type.
|
||||
- Op5's output is converted from u16 to u8. Op6 consumes the u8 type.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
overrides: TensorQuantOverridesHelper,
|
||||
producers: dict[str, onnx.NodeProto],
|
||||
consumers: dict[str, list[onnx.NodeProto]],
|
||||
value_infos: dict[str, onnx.ValueInfoProto],
|
||||
initializers: dict[str, onnx.TensorProto],
|
||||
):
|
||||
"""
|
||||
Params:
|
||||
overrides: The initial tensor quantization overrides to fix.
|
||||
producers: Dictionary that maps a tensor name to the producer node that generates the tensor.
|
||||
consumers: Dictionary that maps a tensor name to the consumer nodes that take the tensor as input.
|
||||
value_infos: Dictionary that maps a tensor name to its onnx.ValueInfoProto.
|
||||
initializers: Dictionary that maps an initializer name to its onnx.TensorProto.
|
||||
"""
|
||||
self.overrides = overrides
|
||||
self.consumers = consumers
|
||||
self.producers = producers
|
||||
self.value_infos = value_infos
|
||||
self.initializers = initializers
|
||||
|
||||
@staticmethod
|
||||
def create_from_model(
|
||||
overrides: TensorQuantOverridesHelper, model: onnx.ModelProto, default_activation_qtype: QuantType
|
||||
) -> MixedPrecisionTensorQuantOverridesFixer:
|
||||
"""
|
||||
Helper function that creates an instance of this class from a loaded ONNX model.
|
||||
|
||||
Params:
|
||||
overrides: The initial tensor quantization overrides to fix.
|
||||
model: Loaded ONNX model
|
||||
default_activation_qtype: The intended default activation quantization type.
|
||||
Used to validate the initial overrides.
|
||||
|
||||
Returns:
|
||||
Initialized MixedPrecisionTensorQuantOverridesFixer object
|
||||
"""
|
||||
model = onnx.shape_inference.infer_shapes(model) # Need to infer shapes to get value_infos
|
||||
|
||||
# Build dictionaries that enable convenient lookups of initializers and value_infos by name.
|
||||
initializers = {initializer.name: initializer for initializer in model.graph.initializer}
|
||||
value_infos = {vi.name: vi for vi in model.graph.value_info}
|
||||
value_infos.update({ot.name: ot for ot in model.graph.output})
|
||||
value_infos.update({it.name: it for it in model.graph.input})
|
||||
|
||||
# Ensure that the user-provided initial overrides are actually valid.
|
||||
valid, err = overrides.is_valid(initializers, set(value_infos), default_activation_qtype)
|
||||
if not valid:
|
||||
pprint_overrides = overrides.pprint_str(indent=4)
|
||||
logging.error(f"Provided invalid tensor quantization overrides:\n{pprint_overrides}")
|
||||
raise ValueError(err)
|
||||
|
||||
consumers = {}
|
||||
producers = {}
|
||||
|
||||
# Build dictionaries that map a tensor name to the consumer or producer nodes.
|
||||
for node in model.graph.node:
|
||||
for input_name in node.input:
|
||||
if input_name:
|
||||
if input_name not in consumers:
|
||||
consumers[input_name] = []
|
||||
|
||||
consumers[input_name].append(node)
|
||||
|
||||
for output_name in node.output:
|
||||
producers[output_name] = node
|
||||
|
||||
return MixedPrecisionTensorQuantOverridesFixer(overrides, producers, consumers, value_infos, initializers)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
default_activation_qtype: QuantType,
|
||||
default_activation_symmetric: bool,
|
||||
):
|
||||
"""
|
||||
Fixes the initial tensor quantization overrides (in-place) for use in mixed-precision QDQ models.
|
||||
|
||||
Params:
|
||||
default_activation_qtype: The intended default activation quantization type.
|
||||
default_activation_symmetric: The intended default symmetry used to quantize activations.
|
||||
"""
|
||||
type_requests = self.get_desired_tensor_types(default_activation_qtype, default_activation_symmetric)
|
||||
|
||||
# Use type requests to "fix" tensor quantization overrides by adding
|
||||
# quantization type conversions where necessary.
|
||||
for tensor_name, type_req in type_requests.items():
|
||||
all_consumers = {node.name for node in self.consumers.get(tensor_name, [])}
|
||||
has_producer_req = type_req.producer is not None
|
||||
has_consumer_req = bool(type_req.consumers)
|
||||
|
||||
# Only producer type: Add conversion back to default activation type
|
||||
if has_producer_req and not has_consumer_req:
|
||||
self._update_converted_tensor(
|
||||
tensor_name, type_req.producer, QuantTypeInfo(default_activation_qtype), all_consumers
|
||||
)
|
||||
# Only consumers
|
||||
elif not has_producer_req and has_consumer_req:
|
||||
prod_type_info = self.overrides.get_node_output_qtype_info(tensor_name, default_activation_qtype)
|
||||
consumer_type_info = type_req.consumers[0]
|
||||
|
||||
if prod_type_info != consumer_type_info:
|
||||
self._update_converted_tensor(
|
||||
tensor_name, prod_type_info, consumer_type_info, type_req.consumers[1]
|
||||
)
|
||||
else:
|
||||
if not self._check_nodes_are_not_convert_consumers(tensor_name, type_req.consumers[1]):
|
||||
raise ValueError(
|
||||
f"Tensor override for '{tensor_name}' converts the type for consumers that need the original type."
|
||||
)
|
||||
# Both producer and consumers
|
||||
elif has_producer_req and has_consumer_req:
|
||||
prod_type_info = type_req.producer
|
||||
consumer_type_info = type_req.consumers[0]
|
||||
|
||||
if prod_type_info != consumer_type_info:
|
||||
self._update_converted_tensor(
|
||||
tensor_name, prod_type_info, consumer_type_info, type_req.consumers[1]
|
||||
)
|
||||
else:
|
||||
consumers_for_original_type = all_consumers.difference(type_req.consumers[1])
|
||||
|
||||
if len(consumers_for_original_type) == 0:
|
||||
# All consumers want the overridden type, so no need for convert nodes!
|
||||
# Just add the override to the new new if not already present.
|
||||
if tensor_name not in self.overrides:
|
||||
self.overrides[tensor_name] = [{}]
|
||||
prod_type_info.save_to_dict(self.overrides[tensor_name][0])
|
||||
|
||||
assert "convert" not in self.overrides[tensor_name][0]
|
||||
else:
|
||||
# Some consumers don't want the overridden type.
|
||||
self._update_converted_tensor(
|
||||
tensor_name,
|
||||
prod_type_info,
|
||||
QuantTypeInfo(default_activation_qtype),
|
||||
consumers_for_original_type,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"TypeRequest for tensor {tensor_name} has no producer or consumers.")
|
||||
|
||||
# Done. Check if the overrides are valid.
|
||||
valid, err = self.overrides.is_valid(self.initializers, set(self.value_infos), default_activation_qtype)
|
||||
if not valid:
|
||||
pprint_overrides = self.overrides.pprint_str(indent=4)
|
||||
logging.error(
|
||||
f"Generated invalid tensor quantization overrides for mixed-precision QDQ model:\n{pprint_overrides}"
|
||||
)
|
||||
raise ValueError(err)
|
||||
|
||||
def get_desired_tensor_types(
|
||||
self,
|
||||
default_activation_qtype: QuantType,
|
||||
default_activation_symmetric: bool,
|
||||
) -> dict[str, TensorTypeRequest]:
|
||||
"""
|
||||
Iterates through the initial tensor quantization overrides and builds a set of TensorTypeRequests objects
|
||||
that describe the quantization types required at each tensor. These TensorTypeRequests objects are ultimately
|
||||
used to generated the "fixed" overrides.
|
||||
|
||||
Params:
|
||||
default_activation_qtype: The intended default activation quantization type.
|
||||
default_activation_symmetric: The intended default symmetry used to quantize activations.
|
||||
|
||||
Returns:
|
||||
TensorTypeRequest objects as a dict that maps a tensor name to its requested types.
|
||||
"""
|
||||
type_requests = {}
|
||||
default_activation_type_info = QuantTypeInfo(default_activation_qtype, default_activation_symmetric)
|
||||
|
||||
# Scan tensor overrides for type conversion requests.
|
||||
for tensor_name, override_list in self.overrides.items():
|
||||
if not self.__is_tensor_quantizable(tensor_name):
|
||||
continue # Skip non-quantizable tensors (e.g., not a float)
|
||||
|
||||
if tensor_name in self.initializers:
|
||||
continue # Skip initializers
|
||||
|
||||
if not override_list or len(override_list) > 1:
|
||||
continue # Skip per-channel stuff
|
||||
|
||||
override_dict = override_list[0]
|
||||
quant_type_info = QuantTypeInfo.load_from_dict(override_dict, default_activation_type_info.quant_type)
|
||||
producer_node = self.producers.get(tensor_name) # None if this is a model input
|
||||
|
||||
if quant_type_info != default_activation_type_info and "convert" not in override_dict:
|
||||
if producer_node is not None:
|
||||
self._add_type_requests_for_node(type_requests, quant_type_info, producer_node)
|
||||
|
||||
# Find all consumer nodes of `tensor_name` and update their inputs/outputs to the new type.
|
||||
for consumer_node in self.consumers.get(tensor_name, []):
|
||||
self._add_type_requests_for_node(type_requests, quant_type_info, consumer_node)
|
||||
|
||||
return type_requests
|
||||
|
||||
def _add_type_requests_for_node(
|
||||
self,
|
||||
type_requests: dict[str, TensorTypeRequest],
|
||||
quant_type_info: QuantTypeInfo,
|
||||
node: onnx.NodeProto,
|
||||
):
|
||||
"""
|
||||
Adds TensorTypeRequest objects for a given node, assuming that we want all its inputs and outputs
|
||||
to have the same quantization type (as specified by the `quant_type_info` parameter).
|
||||
|
||||
Params:
|
||||
type_requests: Dictionary of type requests to append to for this node.
|
||||
quant_type_info: The quantization type to use for inputs and outputs.
|
||||
node: The node for which the TensorTypeRequest objects are created and added to type_requests.
|
||||
"""
|
||||
# Add output side
|
||||
for output_name in node.output:
|
||||
if not self.__is_tensor_quantizable(output_name):
|
||||
continue
|
||||
|
||||
if output_name not in type_requests:
|
||||
type_requests[output_name] = TensorTypeRequest(quant_type_info, None)
|
||||
else:
|
||||
if (
|
||||
type_requests[output_name].producer is not None
|
||||
and type_requests[output_name].producer != quant_type_info
|
||||
):
|
||||
raise ValueError(f"Tensor {output_name} has multiple types.")
|
||||
|
||||
type_requests[output_name].producer = quant_type_info
|
||||
|
||||
# Add the consumer side
|
||||
for input_name in node.input:
|
||||
if input_name and input_name not in self.initializers and self.__is_tensor_quantizable(input_name):
|
||||
if input_name not in type_requests:
|
||||
type_requests[input_name] = TensorTypeRequest(None, None)
|
||||
|
||||
if type_requests[input_name].consumers is None:
|
||||
type_requests[input_name].consumers = (quant_type_info, set())
|
||||
|
||||
if type_requests[input_name].consumers[0] != quant_type_info:
|
||||
raise ValueError(f"Tensor {input_name} has consumers requesting different types.")
|
||||
|
||||
if not node.name:
|
||||
raise ValueError(
|
||||
f"Node of type {node.op_type} with output 0 {node.output[0]} does not have a name!"
|
||||
)
|
||||
|
||||
type_requests[input_name].consumers[1].add(node.name)
|
||||
|
||||
def _update_converted_tensor(
|
||||
self,
|
||||
tensor_name: str,
|
||||
producer_type_info: QuantTypeInfo,
|
||||
consumer_type_info: QuantTypeInfo,
|
||||
consumer_names: set[str],
|
||||
):
|
||||
"""
|
||||
Updates the tensor quantization overrides for a tensor that is converted from one type to another.
|
||||
|
||||
Params:
|
||||
tensor_name: The name of the tensor for which to update overrides.
|
||||
producer_type_info: Info for the tensor's produced type.
|
||||
consumer_type_info: Info for the tensor's consumed (i.e., converted) type.
|
||||
consumer_names: Nodes names of consumers that consume the converted type.
|
||||
"""
|
||||
if tensor_name not in self.overrides or not self.overrides[tensor_name]:
|
||||
self.overrides[tensor_name] = [{}]
|
||||
producer_type_info.save_to_dict(self.overrides[tensor_name][0])
|
||||
|
||||
overrides = self.overrides[tensor_name][0]
|
||||
if producer_type_info != QuantTypeInfo.load_from_dict(overrides):
|
||||
raise ValueError(f"Desired producer quant_type for {tensor_name} doesn't match existing type.")
|
||||
|
||||
if consumer_names:
|
||||
if "convert" not in overrides:
|
||||
overrides["convert"] = {}
|
||||
consumer_type_info.save_to_dict(overrides["convert"])
|
||||
|
||||
convert_dict = overrides["convert"]
|
||||
if consumer_type_info != QuantTypeInfo.load_from_dict(convert_dict):
|
||||
raise ValueError(f"Desired consumer quant_type for {tensor_name} doesn't match existing type.")
|
||||
|
||||
if "recv_nodes" not in convert_dict:
|
||||
convert_dict["recv_nodes"] = set()
|
||||
|
||||
convert_dict["recv_nodes"].update(consumer_names)
|
||||
|
||||
def _check_nodes_are_not_convert_consumers(self, tensor_name: str, node_names: set[str]):
|
||||
"""
|
||||
Returns true if the given nodes do not consume/receive a converted quantization type.
|
||||
|
||||
Params:
|
||||
tensor_name: The name of the tensor to check.
|
||||
node_names: Set of node names that should not be consumers of the converted type.
|
||||
"""
|
||||
if tensor_name not in self.overrides or not self.overrides[tensor_name]:
|
||||
return True
|
||||
|
||||
overrides = self.overrides[tensor_name][0]
|
||||
|
||||
if "convert" not in overrides:
|
||||
return True
|
||||
|
||||
convert_dict = overrides["convert"]
|
||||
|
||||
if "recv_nodes" not in convert_dict:
|
||||
return False
|
||||
|
||||
return not convert_dict["recv_nodes"].intersection(node_names)
|
||||
|
||||
def __is_tensor_quantizable(self, tensor_name):
|
||||
weight = self.initializers.get(tensor_name)
|
||||
if weight is not None:
|
||||
if weight.data_type in (onnx.TensorProto.FLOAT, onnx.TensorProto.FLOAT16):
|
||||
return True
|
||||
elif tensor_name in self.value_infos:
|
||||
vi = self.value_infos[tensor_name]
|
||||
if vi.type.HasField("tensor_type") and vi.type.tensor_type.elem_type in (
|
||||
onnx.TensorProto.FLOAT,
|
||||
onnx.TensorProto.FLOAT16,
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -0,0 +1,335 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
import onnx
|
||||
|
||||
from ....tools.onnx_model_utils import fix_output_shapes, make_input_shape_fixed
|
||||
from ....tools.remove_initializer_from_input import remove_initializer_from_input
|
||||
from ...fusions import FusionGelu, FusionLayerNormalization
|
||||
from ...onnx_model import ONNXModel
|
||||
from ...quant_utils import save_and_reload_model_with_shape_infer
|
||||
from .fusion_lpnorm import FusionLpNormalization
|
||||
from .fusion_spacetodepth import FusionSpaceToDepth
|
||||
|
||||
|
||||
def qnn_preprocess_model(
|
||||
model_input: str | Path | onnx.ModelProto,
|
||||
model_output: str | Path,
|
||||
exclude_initializer_from_input: bool = False,
|
||||
fuse_layernorm: bool = False,
|
||||
save_as_external_data: bool = False,
|
||||
all_tensors_to_one_file: bool = False,
|
||||
external_data_location: str | None = None,
|
||||
external_data_size_threshold: int = 1024,
|
||||
external_data_convert_attribute: bool = False,
|
||||
inputs_to_make_channel_last: list[str] | None = None,
|
||||
outputs_to_make_channel_last: list[str] | None = None,
|
||||
dynamic_input_shapes: list[tuple[str, str]] | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
If necessary, this method creates a new "pre-processed" model in preparation for
|
||||
quantization of a model to be used in QNN EP. Returns true if a new model was created.
|
||||
|
||||
This method perfoms the following operations:
|
||||
- Fuse Erf sequence into a single Gelu node.
|
||||
- Fuse ReduceL2 sequence into a single LpNormalization node (p == 2).
|
||||
- (Optional) Fuse ReduceMean sequence into a single LayerNormalization node.
|
||||
|
||||
Args:
|
||||
model_input: Path to the input model file or ModelProto.
|
||||
model_output: Path the output model file, which is only created if this method returns True.
|
||||
exclude_initializer_from_input: A bool specifying whether to exclude initializer from input.
|
||||
Defaults to False.
|
||||
fuse_layernorm: True if ReduceMean sequences should be fused into LayerNormalization nodes.
|
||||
Defaults to False.
|
||||
save_as_external_data: True if output model should be saved with external data. Defaults to false.
|
||||
all_tensors_to_one_file: Effective only if save_as_external_data is true. Defaults to false.
|
||||
If true, save all tensors to one external file specified by external_data_location.
|
||||
If false, save each tensor to a file named with the tensor name.
|
||||
external_data_location: Effective only if save_as_external_data is true. Defaults to None.
|
||||
Specify the external file to which all tensors are saved. Path is relative
|
||||
to the model path. If not specified, the model's name is used.
|
||||
external_data_size_threshold: Effective only if save_as_external_data is true. Defaults to 1024.
|
||||
Tensors with a data size >= external_data_size_threshold are converted to external data.
|
||||
To convert every tensor with raw data to external data, set to 0.
|
||||
external_data_convert_attribute: Effective only if save_as_external_data is true. Defaults to false.
|
||||
If true, convert all tensors to external data.
|
||||
If false, convert only non-attribute tensors to external data.
|
||||
inputs_to_make_channel_last: List of graph input names to transpose to be "channel-last". For example,
|
||||
if "input0" originally has the shape (N, C, D1, D2, ..., Dn), the resulting model will change input0's
|
||||
shape to (N, D1, D2, ..., Dn, C) and add a transpose node after it.
|
||||
|
||||
Original:
|
||||
input0 (N, C, D1, D2, ..., Dn) --> <Nodes>
|
||||
|
||||
Updated:
|
||||
input0 (N, D1, D2, ..., Dn, C) --> Transpose --> input0_chanfirst (N, C, D1, D2, ..., Dn) --> <Nodes>
|
||||
|
||||
This can potentially improve inference latency for QDQ models running on QNN EP because the
|
||||
additional transpose node may allow other transpose nodes inserted during ORT layout transformation
|
||||
to cancel out.
|
||||
outputs_to_make_channel_last: List of graph output names to transpose to be "channel-last". For example,
|
||||
if "output0" originally has the shape (N, C, D1, D2, ..., Dn), the resulting model will change output0's
|
||||
shape to (N, D1, D2, ..., Dn, C) and add a transpose node before it.
|
||||
|
||||
Original:
|
||||
<Nodes> --> output0 (N, C, D1, D2, ..., Dn)
|
||||
|
||||
Updated:
|
||||
<Nodes> --> output0_chanfirst (N, C, D1, D2, ..., Dn) --> Transpose --> output0 (N, D1, D2, ..., Dn, C)
|
||||
|
||||
This can potentially improve inference latency for QDQ models running on QNN EP because the
|
||||
additional transpose node may allow other transpose nodes inserted during ORT layout transformation
|
||||
to cancel out.
|
||||
dynamic_input_shapes: A list of tuples specifying model input name to and its static shape in comma seprated
|
||||
format, for example: [('input', '1,3,256,256')]. Defaults to None.
|
||||
"""
|
||||
modified = False
|
||||
model = model_input if isinstance(model_input, onnx.ModelProto) else onnx.load_model(model_input)
|
||||
model = save_and_reload_model_with_shape_infer(model)
|
||||
onnx_model = ONNXModel(model)
|
||||
|
||||
# Optionally, fix the dynamic input shapes.
|
||||
if dynamic_input_shapes:
|
||||
for input_name, input_shape_str in dynamic_input_shapes:
|
||||
input_shape = [int(i) for i in input_shape_str.split(",")]
|
||||
make_input_shape_fixed(onnx_model.graph(), input_name, input_shape)
|
||||
fix_output_shapes(onnx_model.model)
|
||||
modified = True
|
||||
|
||||
# Exclude initializer from input if model.ir_version >= 4
|
||||
if exclude_initializer_from_input:
|
||||
modified |= remove_initializer_from_input(onnx_model.model)
|
||||
|
||||
# Fuse Erf sequence into a single Gelu
|
||||
fusion_gelu = FusionGelu(onnx_model)
|
||||
if fusion_gelu.apply():
|
||||
modified = True
|
||||
|
||||
# Fuse ReduceL2 sequence into a single LpNormalization node with p == 2.
|
||||
fusion_lpnorm = FusionLpNormalization(onnx_model)
|
||||
if fusion_lpnorm.apply():
|
||||
modified = True
|
||||
|
||||
# Fuse Reshape/Transpose sequence into a single SpaceToDepth.
|
||||
fusion_s2d = FusionSpaceToDepth(onnx_model)
|
||||
if fusion_s2d.apply():
|
||||
modified = True
|
||||
|
||||
# Optionally, fuse ReduceMean sequence into a single LayerNormalization node.
|
||||
if fuse_layernorm:
|
||||
onnx_opset = next(x for x in model.opset_import if x.domain == "" or x.domain == "ai.onnx")
|
||||
|
||||
# Need opset >= 17 to use LayerNormalization.
|
||||
if onnx_opset.version < 17:
|
||||
logging.warning(
|
||||
"Unable to fuse ReduceMean sequence into a LayerNormalization node. "
|
||||
"ONNX model must use an opset >= 17 in order to use LayerNormalization, "
|
||||
f"but found version {onnx_opset.version}. Please use onnx.version_converter to update your model."
|
||||
)
|
||||
else:
|
||||
fusion_layernorm = FusionLayerNormalization(onnx_model)
|
||||
if fusion_layernorm.apply():
|
||||
modified = True
|
||||
|
||||
# Optionally, transpose inputs and/or outputs to make them "channel-last".
|
||||
if inputs_to_make_channel_last or outputs_to_make_channel_last:
|
||||
transpose_node_prefix = "Transpose_channel_"
|
||||
transpose_node_suffix: int = onnx_model.get_largest_node_name_suffix(transpose_node_prefix) + 1
|
||||
update_io_to_channel_last(
|
||||
onnx_model.model,
|
||||
inputs_to_make_channel_last,
|
||||
outputs_to_make_channel_last,
|
||||
transpose_node_name_prefix=transpose_node_prefix,
|
||||
transpose_node_name_start_suffix=transpose_node_suffix,
|
||||
)
|
||||
modified = True
|
||||
|
||||
# Make sure all nodes have a name.
|
||||
unnamed_node_prefix = "qnn_preproc_node_"
|
||||
available_suffix = onnx_model.get_largest_node_name_suffix(unnamed_node_prefix) + 1
|
||||
for node in onnx_model.model.graph.node:
|
||||
if node.op_type != "Constant" and not node.name:
|
||||
new_node_name = f"{unnamed_node_prefix}{available_suffix!s}"
|
||||
available_suffix += 1
|
||||
node.name = new_node_name
|
||||
modified = True
|
||||
logging.warning(f"Node of type {node.op_type} does not have a name. Renamed to {new_node_name}.")
|
||||
|
||||
if modified:
|
||||
onnx_model.topological_sort()
|
||||
onnx.save_model(
|
||||
model,
|
||||
model_output,
|
||||
save_as_external_data=save_as_external_data,
|
||||
all_tensors_to_one_file=all_tensors_to_one_file,
|
||||
location=external_data_location,
|
||||
size_threshold=external_data_size_threshold,
|
||||
convert_attribute=external_data_convert_attribute,
|
||||
)
|
||||
|
||||
return modified
|
||||
|
||||
|
||||
class InputOutputNameMap:
|
||||
def __init__(
|
||||
self,
|
||||
orig_tensor_names: set[str],
|
||||
orig_graph_inputs: dict[str, onnx.ValueInfoProto],
|
||||
orig_graph_outputs: dict[str, onnx.ValueInfoProto],
|
||||
):
|
||||
self.orig_tensor_names = orig_tensor_names
|
||||
self.orig_graph_inputs = orig_graph_inputs
|
||||
self.orig_graph_outputs = orig_graph_outputs
|
||||
self.updated_io_names = {}
|
||||
self.new_value_infos = []
|
||||
|
||||
def get_new_name(self, orig_name: str):
|
||||
if orig_name in self.updated_io_names:
|
||||
return self.updated_io_names[orig_name]
|
||||
|
||||
# Make a new tensor name that is unique among all tensors in the graph.
|
||||
prefix: str = f"{orig_name}_channel_first_"
|
||||
suffix: int = -1
|
||||
for tensor_name in self.orig_tensor_names:
|
||||
if tensor_name.startswith(prefix) and tensor_name[len(prefix) :].isdigit():
|
||||
index = int(tensor_name[len(prefix) :])
|
||||
suffix = max(suffix, index)
|
||||
|
||||
suffix += 1 # This is the first available suffix.
|
||||
new_name = f"{prefix}{suffix!s}"
|
||||
|
||||
# Add new value_info objects for these new tensors.
|
||||
orig_value_info = self.orig_graph_inputs.get(orig_name) or self.orig_graph_outputs[orig_name]
|
||||
value_info_proto = onnx.ValueInfoProto()
|
||||
value_info_proto.CopyFrom(orig_value_info)
|
||||
value_info_proto.name = new_name
|
||||
self.new_value_infos.append(value_info_proto)
|
||||
|
||||
self.updated_io_names[orig_name] = new_name
|
||||
return self.updated_io_names[orig_name]
|
||||
|
||||
|
||||
def update_io_to_channel_last(
|
||||
model: onnx.ModelProto,
|
||||
inputs_to_update: list[str] | None,
|
||||
outputs_to_update: list[str] | None,
|
||||
transpose_node_name_prefix: str = "Transpose_channel_",
|
||||
transpose_node_name_start_suffix: int = 0,
|
||||
):
|
||||
inputs_to_update = set(inputs_to_update or [])
|
||||
outputs_to_update = set(outputs_to_update or [])
|
||||
|
||||
if not inputs_to_update and not outputs_to_update:
|
||||
return
|
||||
|
||||
graph = model.graph
|
||||
orig_graph_inputs = {ginput.name: ginput for ginput in graph.input}
|
||||
orig_graph_outputs = {goutput.name: goutput for goutput in graph.output}
|
||||
|
||||
# Check that the user passed in actual input and output names.
|
||||
for input_name in inputs_to_update:
|
||||
if input_name not in orig_graph_inputs:
|
||||
raise ValueError(f"{input_name} is not a graph input")
|
||||
|
||||
for output_name in outputs_to_update:
|
||||
if output_name not in orig_graph_outputs:
|
||||
raise ValueError(f"{output_name} is not a graph output")
|
||||
|
||||
orig_tensor_names = set()
|
||||
orig_tensor_names.update(set(orig_graph_inputs))
|
||||
orig_tensor_names.update(set(orig_graph_outputs))
|
||||
orig_tensor_names.update(input_name for node in graph.node for input_name in node.input if input_name)
|
||||
|
||||
# Maps original input (or output) name to its updated name used within the graph.
|
||||
io_map = InputOutputNameMap(orig_tensor_names, orig_graph_inputs, orig_graph_outputs)
|
||||
|
||||
# Update each node's inputs/outputs to use the transposed versions.
|
||||
for node in graph.node:
|
||||
for i in range(len(node.input)):
|
||||
if node.input[i] and node.input[i] in inputs_to_update:
|
||||
node.input[i] = io_map.get_new_name(node.input[i])
|
||||
elif node.input[i] and node.input[i] in outputs_to_update:
|
||||
node.input[i] = io_map.get_new_name(node.input[i])
|
||||
|
||||
for i in range(len(node.output)):
|
||||
if node.output[i] in outputs_to_update:
|
||||
node.output[i] = io_map.get_new_name(node.output[i])
|
||||
|
||||
# Update graph inputs to channel-last and a Transpose (to channel-first) after each.
|
||||
for g_input_name in inputs_to_update:
|
||||
g_input = orig_graph_inputs[g_input_name]
|
||||
|
||||
if not g_input.type.HasField("tensor_type") or not g_input.type.tensor_type.HasField("shape"):
|
||||
raise ValueError(f"Expected input {g_input.name} to have a tensor_type with a shape")
|
||||
|
||||
input_shape = g_input.type.tensor_type.shape
|
||||
input_rank = len(input_shape.dim)
|
||||
|
||||
if input_rank < 3:
|
||||
raise ValueError(f"Expected input {g_input.name} to be of rank >= 3")
|
||||
|
||||
channel_dim = onnx.TensorShapeProto.Dimension()
|
||||
channel_dim.CopyFrom(input_shape.dim[1])
|
||||
for i in range(1, input_rank - 1):
|
||||
input_shape.dim[i].CopyFrom(input_shape.dim[i + 1])
|
||||
input_shape.dim[input_rank - 1].CopyFrom(channel_dim)
|
||||
|
||||
transpose_perm = list(range(input_rank))
|
||||
for i in range(input_rank):
|
||||
transpose_perm[i] = i if i < 1 else i - 1
|
||||
transpose_perm[1] = input_rank - 1
|
||||
|
||||
transpose_node = onnx.helper.make_node(
|
||||
"Transpose",
|
||||
name=f"{transpose_node_name_prefix}{transpose_node_name_start_suffix!s}",
|
||||
inputs=[g_input.name],
|
||||
outputs=[io_map.get_new_name(g_input.name)],
|
||||
perm=transpose_perm,
|
||||
)
|
||||
transpose_node_name_start_suffix += 1
|
||||
|
||||
graph.node.extend([transpose_node])
|
||||
|
||||
# Update graph outputs to channel-last and a Transpose (from channel-first) before each.
|
||||
for g_output_name in outputs_to_update:
|
||||
g_output = orig_graph_outputs[g_output_name]
|
||||
if not g_output.type.HasField("tensor_type") or not g_output.type.tensor_type.HasField("shape"):
|
||||
raise ValueError(f"Expected output {g_output.name} to have a tensor_type with a shape")
|
||||
|
||||
output_shape = g_output.type.tensor_type.shape
|
||||
output_rank = len(output_shape.dim)
|
||||
|
||||
if output_rank < 3:
|
||||
raise ValueError(f"Expected output {g_output.name} to be of rank >= 3")
|
||||
|
||||
channel_dim = onnx.TensorShapeProto.Dimension()
|
||||
channel_dim.CopyFrom(output_shape.dim[1])
|
||||
for i in range(1, output_rank - 1):
|
||||
output_shape.dim[i].CopyFrom(output_shape.dim[i + 1])
|
||||
output_shape.dim[output_rank - 1].CopyFrom(channel_dim)
|
||||
|
||||
transpose_perm = list(range(output_rank))
|
||||
for i in range(output_rank):
|
||||
transpose_perm[i] = i if i == 0 else i + 1
|
||||
transpose_perm[output_rank - 1] = 1
|
||||
|
||||
transpose_node = onnx.helper.make_node(
|
||||
"Transpose",
|
||||
name=f"{transpose_node_name_prefix}{transpose_node_name_start_suffix!s}",
|
||||
inputs=[io_map.get_new_name(g_output.name)],
|
||||
outputs=[g_output.name],
|
||||
perm=transpose_perm,
|
||||
)
|
||||
transpose_node_name_start_suffix += 1
|
||||
|
||||
graph.node.extend([transpose_node])
|
||||
|
||||
graph.value_info.extend(io_map.new_value_infos)
|
||||
@@ -0,0 +1,406 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
|
||||
from ...calibrate import CalibrationDataReader, CalibrationMethod
|
||||
from ...quant_utils import QuantType
|
||||
from ...quantize import StaticQuantConfig
|
||||
from ...tensor_quant_overrides import TensorQuantOverridesHelper
|
||||
from .mixed_precision_overrides_utils import MixedPrecisionTensorQuantOverridesFixer
|
||||
|
||||
Q16_TYPES = {QuantType.QInt16, QuantType.QUInt16}
|
||||
Q8_TYPES = {QuantType.QInt8, QuantType.QUInt8}
|
||||
Q4_TYPES = {QuantType.QInt4, QuantType.QUInt4}
|
||||
OP_TYPES_TO_EXCLUDE = {"Cast"}
|
||||
MODEL_SIZE_THRESHOLD = 2147483648 # Quant model should use external data if >= 2GB
|
||||
|
||||
|
||||
def warn_unable_to_override(
|
||||
node: onnx.NodeProto,
|
||||
what_str: str,
|
||||
tensor_name: str,
|
||||
io_kind: str,
|
||||
):
|
||||
logging.warning(
|
||||
f"Unable to override {what_str} for {node.op_type} node's {io_kind} "
|
||||
"because it has already been overridden! Check the initial quantization overrides provided "
|
||||
"to get_qnn_qdq_config() if the generated QDQ model does not run on QNN EP. "
|
||||
f"Node name: {node.name}, {io_kind} name: {tensor_name}"
|
||||
)
|
||||
|
||||
|
||||
def get_qnn_qdq_config(
|
||||
model_input: str | Path | onnx.ModelProto,
|
||||
calibration_data_reader: CalibrationDataReader,
|
||||
calibrate_method: CalibrationMethod = CalibrationMethod.MinMax,
|
||||
activation_type: QuantType = QuantType.QUInt8,
|
||||
weight_type: QuantType = QuantType.QUInt8,
|
||||
per_channel: bool = False,
|
||||
init_overrides: dict[str, list[dict[str, Any]]] | None = None,
|
||||
add_qtype_converts: bool = True,
|
||||
activation_symmetric: bool = False,
|
||||
weight_symmetric: bool | None = None,
|
||||
keep_removable_activations: bool = False,
|
||||
stride: int | None = None,
|
||||
calibration_providers: list[str] | None = None,
|
||||
op_types_to_quantize: list[str] | None = None,
|
||||
nodes_to_exclude: list[str] | None = None,
|
||||
) -> StaticQuantConfig:
|
||||
"""
|
||||
Returns a static quantization configuration suitable for running QDQ models on QNN EP.
|
||||
This is done primarily by setting tensor-level quantization overrides.
|
||||
|
||||
Params:
|
||||
model_input: Path to the input model file or ModelProto.
|
||||
calibration_data_reader: Calibration data reader.
|
||||
calibrate_methode: The calibration method. Defaults to MinMax.
|
||||
activation_type: The default activation quantization type. Defaults to QUInt8.
|
||||
weight_type: The default weight quantization type. Defaults to QUInt8.
|
||||
per_channel: Global option that determines if a fixed set of operator types should be quantized per-channel.
|
||||
Defaults to false. Alternatively, use the tensor-level `init_overrides` to select individual operators
|
||||
and their quantization axes.
|
||||
|
||||
If set, the quantization tool uses per-channel quantization for the following operator types and inputs:
|
||||
- Conv:
|
||||
- input[1] on axis 0
|
||||
- input[2] (bias) on axis 0
|
||||
- ConvTranspose:
|
||||
- input[1] on axis 1
|
||||
- input[2] (bias) on axis 0
|
||||
init_overrides: Initial tensor-level quantization overrides. Defaults to None. This function updates of a copy
|
||||
of these overrides with any necessary adjustments and includes them in the returned
|
||||
configuration object (i.e., config.extra_options['TensorQuantOverrides']).
|
||||
|
||||
The key is a tensor name and the value is a list of dictionaries. For per-tensor quantization, the list
|
||||
contains a single dictionary. For per-channel quantization, the list contains either a dictionary for
|
||||
each channel in the tensor or a single dictionary that is assumed to apply to all channels. An 'axis'
|
||||
key must be present in the first dictionary for per-channel quantization.
|
||||
|
||||
Each dictionary contains optional overrides with the following keys and values.
|
||||
'quant_type' = QuantType : The tensor's quantization data type.
|
||||
'axis' = Int : The per-channel axis. Must be present for per-channel weights.
|
||||
'scale' = Float : The scale value to use. Must also specify `zero_point` if set.
|
||||
'zero_point' = Int : The zero-point value to use. Must also specify `scale` is set.
|
||||
'symmetric' = Bool : If the tensor should use symmetric quantization. Invalid if also
|
||||
set `scale` or `zero_point`.
|
||||
'reduce_range' = Bool : If the quantization range should be reduced. Invalid if also
|
||||
set `scale` or `zero_point`. Only valid for initializers.
|
||||
'rmax' = Float : Override the maximum real tensor value in calibration data.
|
||||
Invalid if also set `scale` or `zero_point`.
|
||||
'rmin' = Float : Override the minimum real tensor value in calibration data.
|
||||
Invalid if also set `scale` or `zero_point`.
|
||||
'convert' = Dict : A nested dictionary with the same keys for an activation
|
||||
tensor that should be converted to another quantization type.
|
||||
'convert["recv_nodes"] = Set : Set of node names that consume the converted activation,
|
||||
other nodes get the original type. If not specified,
|
||||
assume all consumer nodes get the converted type.
|
||||
add_qtype_converts: True if this function should automatically add "convert" entries to the provided
|
||||
`init_overrides` to ensure that operators use valid input/output types (activations only).
|
||||
Ex: if you override the output of an Add to 16-bit, this option ensures that the activation inputs
|
||||
of the Add are also up-converted to 16-bit and that data types for surrounding ops are converted
|
||||
appropriately. Refer to the documentation in mixed_precision_overrides_utils.py for additional details.
|
||||
activation_symmetric: True if activations should be quantized symmetrically (i.e, rmax == -rmin) by default.
|
||||
Defaults to false. For int8 and int16, this results in zero-point values of 0. For uint8 and uin16,
|
||||
the zero-point values are 128 and 32,768, respectively.
|
||||
weight_symmetric: True if weights should be quantized symmetrically (i.e., rmax == -rmin) by default.
|
||||
Defaults to None. If set to None, weight_symmetric is assumed true if the weight_type is a signed int.
|
||||
keep_removable_activations: Defaults to false. If true, "removable" activations (e.g., Clip or Relu) will not
|
||||
be removed, and will be explicitly represented in the QDQ model. If false, these activations
|
||||
are automatically removed if activations are asymmetrically quantized. Keeping these activations
|
||||
is necessary if optimizations or EP transformations will later remove
|
||||
QuantizeLinear/DequantizeLinear operators from the model.
|
||||
calibration_providers: Execution providers to run the session during calibration. Default is None which uses
|
||||
[ "CPUExecutionProvider" ].
|
||||
op_types_to_quantize: If set to None, all operator types will be quantized except for OP_TYPES_TO_EXCLUDE
|
||||
nodes_to_exclude: List of nodes names to exclude from quantization. The nodes in this list will be excluded from
|
||||
quantization when it is not None.
|
||||
|
||||
Returns:
|
||||
A StaticQuantConfig object
|
||||
"""
|
||||
if weight_symmetric is None:
|
||||
weight_symmetric = weight_type in {QuantType.QInt8, QuantType.QInt16}
|
||||
|
||||
model = (
|
||||
model_input
|
||||
if isinstance(model_input, onnx.ModelProto)
|
||||
else onnx.load_model(model_input, load_external_data=False)
|
||||
)
|
||||
|
||||
op_types = set()
|
||||
model_has_external_data = False
|
||||
name_to_initializer = {}
|
||||
|
||||
# Build map of initializers (name -> initializer) and
|
||||
# check if the model has external data.
|
||||
for initializer in model.graph.initializer:
|
||||
name_to_initializer[initializer.name] = initializer
|
||||
if onnx.external_data_helper.uses_external_data(initializer):
|
||||
model_has_external_data = True
|
||||
|
||||
overrides_helper = TensorQuantOverridesHelper(copy.deepcopy(init_overrides) if init_overrides else {})
|
||||
|
||||
if not overrides_helper.empty() and add_qtype_converts:
|
||||
# Fix mixed-precision overrides.
|
||||
overrides_fixer = MixedPrecisionTensorQuantOverridesFixer.create_from_model(
|
||||
overrides_helper, model, activation_type
|
||||
)
|
||||
overrides_fixer.apply(activation_type, activation_symmetric)
|
||||
|
||||
# Setup quantization overrides for specific operator types to ensure compatibility with QNN EP.
|
||||
qnn_compat = QnnCompatibilityOverrides(
|
||||
activation_type,
|
||||
weight_type,
|
||||
activation_symmetric,
|
||||
weight_symmetric,
|
||||
per_channel,
|
||||
overrides_helper,
|
||||
name_to_initializer,
|
||||
)
|
||||
|
||||
op_types_to_quantize_set = set(op_types_to_quantize) if op_types_to_quantize else None
|
||||
nodes_to_exclude_set = set(nodes_to_exclude) if nodes_to_exclude else None
|
||||
|
||||
for node in model.graph.node:
|
||||
if op_types_to_quantize_set and node.op_type not in op_types_to_quantize_set:
|
||||
continue
|
||||
if nodes_to_exclude_set and node.name in nodes_to_exclude_set:
|
||||
continue
|
||||
op_types.add(node.op_type)
|
||||
qnn_compat.process_node(node)
|
||||
|
||||
extra_options = {
|
||||
"MinimumRealRange": 0.0001,
|
||||
"DedicatedQDQPair": False, # Let ORT optimizer duplicate DQ nodes
|
||||
"QDQKeepRemovableActivations": keep_removable_activations,
|
||||
"TensorQuantOverrides": overrides_helper.get_dict(),
|
||||
"ActivationSymmetric": activation_symmetric,
|
||||
"WeightSymmetric": weight_symmetric,
|
||||
"CalibStridedMinMax": stride,
|
||||
}
|
||||
|
||||
# ONNX opset < 21 does not support 16-bit quantization, so must use 'com.microsoft' domain
|
||||
# on Q/DQ operators if using 16-bit or 4-bit quantization.
|
||||
onnx_opset = next(x for x in model.opset_import if x.domain == "" or x.domain == "ai.onnx")
|
||||
if onnx_opset.version < 21:
|
||||
opset21_types = Q16_TYPES.union(Q4_TYPES)
|
||||
overrides_have_opset21_types = any(t in opset21_types for t in overrides_helper.get_quant_types())
|
||||
if activation_type in opset21_types or weight_type in opset21_types or overrides_have_opset21_types:
|
||||
extra_options["UseQDQContribOps"] = True
|
||||
|
||||
return StaticQuantConfig(
|
||||
calibration_data_reader,
|
||||
calibrate_method=calibrate_method,
|
||||
activation_type=activation_type,
|
||||
weight_type=weight_type,
|
||||
op_types_to_quantize=(
|
||||
op_types_to_quantize if op_types_to_quantize else list(op_types.difference(OP_TYPES_TO_EXCLUDE))
|
||||
),
|
||||
nodes_to_exclude=nodes_to_exclude,
|
||||
per_channel=per_channel,
|
||||
use_external_data_format=(model_has_external_data or model.ByteSize() >= MODEL_SIZE_THRESHOLD),
|
||||
calibration_providers=calibration_providers,
|
||||
extra_options=extra_options,
|
||||
)
|
||||
|
||||
|
||||
class QnnCompatibilityOverrides:
|
||||
"""
|
||||
Helper that processes nodes to generate quantization overrides that make the resulting QDQ model
|
||||
compatible with QNN EP.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
default_activation_qtype: QuantType,
|
||||
default_weight_qtype: QuantType,
|
||||
activation_symmetric: bool,
|
||||
weight_symmetric: bool,
|
||||
per_channel: bool,
|
||||
overrides: TensorQuantOverridesHelper,
|
||||
initializers: dict[str, onnx.TensorProto],
|
||||
):
|
||||
self.default_activation_qtype = default_activation_qtype
|
||||
self.default_weight_qtype = default_weight_qtype
|
||||
self.activation_symmetric = activation_symmetric
|
||||
self.weight_symmetric = weight_symmetric
|
||||
self.per_channel = per_channel
|
||||
self.overrides = overrides
|
||||
self.initializers = initializers
|
||||
|
||||
self.process_fns = {
|
||||
"MatMul": self._process_matmul,
|
||||
"LayerNormalization": self._process_layernorm,
|
||||
"Sigmoid": self._process_sigmoid,
|
||||
"Tanh": self._process_tanh,
|
||||
}
|
||||
|
||||
def process_node(self, node: onnx.NodeProto):
|
||||
process_fn = self.process_fns.get(node.op_type)
|
||||
|
||||
if process_fn is not None:
|
||||
process_fn(node)
|
||||
|
||||
def _make_static_inputs_use_default_weight_type(self, node: onnx.NodeProto):
|
||||
"""
|
||||
Overrides initializer input(s) to use the default weight type if:
|
||||
- The default weight type is 8-bit
|
||||
- One of the inputs is a 16-bit activation
|
||||
- The other input is an initializer (per-tensor quantized)
|
||||
|
||||
This is necessary because the quantization tool does not assign MatMul or LayerNorm initializer
|
||||
inputs the default weight type. Instead, it assigns the default activation type.
|
||||
"""
|
||||
if self.default_weight_qtype not in Q8_TYPES:
|
||||
return
|
||||
|
||||
input_16bit_act_name = None
|
||||
input_weight_name = None
|
||||
|
||||
# Loop through first 2 inputs to find a 16-bit activation and a (per-tensor) weight.
|
||||
for i in range(2):
|
||||
input_name = node.input[i]
|
||||
if not input_name:
|
||||
continue
|
||||
|
||||
is_weight = input_name in self.initializers
|
||||
qtype_info = self.overrides.get_node_input_qtype_info(
|
||||
input_name,
|
||||
node.name,
|
||||
default_qtype=None if is_weight else self.default_activation_qtype,
|
||||
)
|
||||
|
||||
if qtype_info.axis is not None:
|
||||
return # Don't process MatMul with a per-channel quantized input.
|
||||
|
||||
if (
|
||||
is_weight
|
||||
and qtype_info.quant_type == self.default_weight_qtype
|
||||
and qtype_info.symmetric == self.weight_symmetric
|
||||
):
|
||||
return # Return. Weight is already overridden to use the desired weight type.
|
||||
|
||||
if is_weight:
|
||||
input_weight_name = input_name
|
||||
elif qtype_info.quant_type in Q16_TYPES:
|
||||
input_16bit_act_name = input_name
|
||||
|
||||
# Override initializer input to use the default weight type.
|
||||
if input_16bit_act_name and input_weight_name:
|
||||
did_update = self.overrides.update_tensor_overrides(
|
||||
input_weight_name,
|
||||
{"quant_type": self.default_weight_qtype, "symmetric": self.weight_symmetric},
|
||||
overwrite=False,
|
||||
)
|
||||
|
||||
if not did_update:
|
||||
warn_unable_to_override(node, "quant_type/symmetric", input_weight_name, "input weight")
|
||||
|
||||
def _process_matmul(self, node: onnx.NodeProto):
|
||||
assert node.op_type == "MatMul", f"Expected MatMul, but got {node.op_type}"
|
||||
|
||||
if not self.per_channel:
|
||||
self._make_static_inputs_use_default_weight_type(node)
|
||||
return
|
||||
|
||||
# QNN does not support per-channel MatMul. However, the ORT quantization tool attempts to use per-channel
|
||||
# quantization for MatMul by default *if* the global per_channel setting is enabled. So, we need to
|
||||
# provide explicit per-tensor quantization overrides for MatMul if per_channel is enabled and
|
||||
# the user did not provide any other overrides.
|
||||
for input_name in node.input:
|
||||
is_weight_no_overrides = input_name in self.initializers and input_name not in self.overrides
|
||||
if is_weight_no_overrides:
|
||||
self.overrides.update_tensor_overrides(
|
||||
input_name,
|
||||
{"quant_type": self.default_weight_qtype, "symmetric": self.weight_symmetric},
|
||||
)
|
||||
|
||||
def _process_layernorm(self, node: onnx.NodeProto):
|
||||
assert node.op_type == "LayerNormalization", f"Expected LayerNormalization, but got {node.op_type}"
|
||||
|
||||
if not self.per_channel:
|
||||
self._make_static_inputs_use_default_weight_type(node)
|
||||
return
|
||||
|
||||
has_weight_no_overrides = node.input[1] in self.initializers and node.input[1] not in self.overrides
|
||||
has_bias_no_overrides = (
|
||||
len(node.input) > 2
|
||||
and node.input[2]
|
||||
and node.input[2] in self.initializers
|
||||
and node.input[2] not in self.overrides
|
||||
)
|
||||
|
||||
if has_weight_no_overrides or has_bias_no_overrides:
|
||||
# TODO: Make bias input not per-channel. QNN needs it to be per-tensor, but quantizer
|
||||
# tries to makes it per-channel if the weight is also per-channel.
|
||||
raise ValueError(
|
||||
"get_qnn_qdq_config() does not currently support the global per_channel option with LayerNormalization."
|
||||
" Please try using custom overrides that make bias per-tensor quantized."
|
||||
)
|
||||
|
||||
def _process_sigmoid(self, node: onnx.NodeProto):
|
||||
"""
|
||||
Overrides 16-bit Sigmoid's output scale and zero-point as per QNN requirements.
|
||||
"""
|
||||
assert node.op_type == "Sigmoid", f"Expected Sigmoid, but got {node.op_type}"
|
||||
output_type = self.overrides.get_node_output_qtype_info(
|
||||
node.output[0], self.default_activation_qtype
|
||||
).quant_type
|
||||
|
||||
if output_type == QuantType.QUInt16:
|
||||
self.overrides.update_tensor_overrides(
|
||||
node.output[0],
|
||||
{
|
||||
"quant_type": output_type,
|
||||
"scale": np.array(1.0 / 65536.0, dtype=np.float32),
|
||||
"zero_point": np.array(0, dtype=np.uint16),
|
||||
},
|
||||
)
|
||||
elif output_type == QuantType.QInt16:
|
||||
self.overrides.update_tensor_overrides(
|
||||
node.output[0],
|
||||
{
|
||||
"quant_type": output_type,
|
||||
"scale": np.array(1.0 / 32768.0, dtype=np.float32),
|
||||
"zero_point": np.array(0, dtype=np.int16),
|
||||
},
|
||||
)
|
||||
|
||||
def _process_tanh(self, node: onnx.NodeProto):
|
||||
"""
|
||||
Overrides 16-bit Tanh's output scale and zero-point as per QNN requirements.
|
||||
"""
|
||||
assert node.op_type == "Tanh", f"Expected Tanh, but got {node.op_type}"
|
||||
output_type = self.overrides.get_node_output_qtype_info(
|
||||
node.output[0], self.default_activation_qtype
|
||||
).quant_type
|
||||
|
||||
if output_type == QuantType.QUInt16:
|
||||
self.overrides.update_tensor_overrides(
|
||||
node.output[0],
|
||||
{
|
||||
"quant_type": output_type,
|
||||
"scale": np.array(1.0 / 32768.0, dtype=np.float32),
|
||||
"zero_point": np.array(32768, dtype=np.uint16),
|
||||
},
|
||||
)
|
||||
elif output_type == QuantType.QInt16:
|
||||
self.overrides.update_tensor_overrides(
|
||||
node.output[0],
|
||||
{
|
||||
"quant_type": output_type,
|
||||
"scale": np.array(1.0 / 32768.0, dtype=np.float32),
|
||||
"zero_point": np.array(0, dtype=np.int16),
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,4 @@
|
||||
from .fusion import Fusion # noqa: F401
|
||||
from .fusion_gelu import FusionGelu # noqa: F401
|
||||
from .fusion_layernorm import FusionLayerNormalization # noqa: F401
|
||||
from .replace_upsample_with_resize import ReplaceUpsampleWithResize # noqa: F401
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,311 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import deque
|
||||
|
||||
import onnx
|
||||
|
||||
from ..onnx_model import ONNXModel
|
||||
|
||||
|
||||
class Fusion:
|
||||
"""
|
||||
Base class for fusions.
|
||||
"""
|
||||
|
||||
def __init__(self, model: ONNXModel, fused_op_type: str, search_op_type: str):
|
||||
self.search_op_type: str = search_op_type
|
||||
self.fused_op_type: str = fused_op_type
|
||||
self.model: ONNXModel = model
|
||||
self.nodes_to_remove: list = []
|
||||
self.nodes_to_add: list = []
|
||||
|
||||
self._new_node_name_prefix = self.fused_op_type + "_fused_" + self.search_op_type + "_"
|
||||
self._new_node_name_suffix = None # int|None used to create unique node names for the fused ops.
|
||||
|
||||
def fuse(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
):
|
||||
"""
|
||||
Interface function for derived fusion classes. Tries to fuse a node sequence containing
|
||||
the specified node.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def apply(self) -> bool:
|
||||
"""
|
||||
Apply graph fusion on the entire model graph.
|
||||
"""
|
||||
input_name_to_nodes = self.model.input_name_to_nodes()
|
||||
output_name_to_node = self.model.output_name_to_node()
|
||||
|
||||
for node in self.model.nodes():
|
||||
if node.op_type == self.search_op_type:
|
||||
self.fuse(node, input_name_to_nodes, output_name_to_node)
|
||||
|
||||
self.model.remove_nodes(self.nodes_to_remove)
|
||||
self.model.add_nodes(self.nodes_to_add)
|
||||
|
||||
graph_updated = bool(self.nodes_to_remove or self.nodes_to_add)
|
||||
|
||||
if graph_updated:
|
||||
self.model.remove_unused_constant()
|
||||
|
||||
return graph_updated
|
||||
|
||||
def create_unique_node_name(self):
|
||||
prefix = self._new_node_name_prefix
|
||||
|
||||
if self._new_node_name_suffix is None:
|
||||
largest_suffix: int = self.model.get_largest_node_name_suffix(prefix)
|
||||
self._new_node_name_suffix = largest_suffix + 1
|
||||
|
||||
new_name = f"{prefix}{self._new_node_name_suffix!s}"
|
||||
self._new_node_name_suffix += 1
|
||||
|
||||
return new_name
|
||||
|
||||
@staticmethod
|
||||
def is_safe_to_fuse_nodes(
|
||||
nodes_to_remove: list[onnx.NodeProto],
|
||||
keep_outputs: list[str],
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
) -> bool:
|
||||
for node_to_remove in nodes_to_remove:
|
||||
for output_to_remove in node_to_remove.output:
|
||||
if output_to_remove in keep_outputs:
|
||||
continue
|
||||
|
||||
if output_to_remove in input_name_to_nodes:
|
||||
for impacted_node in input_name_to_nodes[output_to_remove]:
|
||||
if impacted_node not in nodes_to_remove:
|
||||
# Not safe to remove nodes since output is used by impacted_node
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def get_node_attribute(node: onnx.NodeProto, attribute_name: str):
|
||||
for attr in node.attribute:
|
||||
if attr.name == attribute_name:
|
||||
value = onnx.helper.get_attribute_value(attr)
|
||||
return value
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def input_index(node_output: str, child_node: onnx.NodeProto) -> int:
|
||||
for index, input_name in enumerate(child_node.input):
|
||||
if input_name == node_output:
|
||||
return index
|
||||
return -1
|
||||
|
||||
@staticmethod
|
||||
def tensor_shape_to_list(tensor_type) -> list[int]:
|
||||
shape_list = []
|
||||
for d in tensor_type.shape.dim:
|
||||
if d.HasField("dim_value"):
|
||||
shape_list.append(d.dim_value) # known dimension
|
||||
elif d.HasField("dim_param"):
|
||||
shape_list.append(d.dim_param) # unknown dimension with symbolic name
|
||||
else:
|
||||
shape_list.append("?") # shall not happen
|
||||
return shape_list
|
||||
|
||||
def get_constant_input(self, node: onnx.NodeProto):
|
||||
for i, inp in enumerate(node.input):
|
||||
value = self.model.get_constant_value(inp)
|
||||
if value is not None:
|
||||
return i, value
|
||||
|
||||
return None, None
|
||||
|
||||
def find_constant_input(self, node: onnx.NodeProto, expected_value: float, delta: float = 0.000001) -> int:
|
||||
i, value = self.get_constant_input(node)
|
||||
if value is not None and value.size == 1 and abs(value - expected_value) < delta:
|
||||
return i
|
||||
|
||||
return -1
|
||||
|
||||
def has_constant_input(self, node: onnx.NodeProto, expected_value: float, delta: float = 0.000001) -> bool:
|
||||
return self.find_constant_input(node, expected_value, delta) >= 0
|
||||
|
||||
def is_constant_with_specified_rank(self, output_name: str, rank: int) -> bool:
|
||||
value = self.model.get_constant_value(output_name)
|
||||
if value is None:
|
||||
return False # Not an initializer
|
||||
|
||||
if len(value.shape) != rank:
|
||||
return False # Wrong dimensions
|
||||
|
||||
return True
|
||||
|
||||
def match_first_parent(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
parent_op_type: str,
|
||||
output_name_to_node: dict[str, onnx.NodeProto] | None = None,
|
||||
exclude: list[onnx.NodeProto] = [], # noqa: B006
|
||||
) -> tuple[onnx.NodeProto | None, int | None]:
|
||||
"""
|
||||
Find parent node based on constraints on op_type.
|
||||
|
||||
Args:
|
||||
node: current node.
|
||||
parent_op_type (str): constraint of parent node op_type.
|
||||
output_name_to_node (dict): dictionary with output name as key, and node as value.
|
||||
exclude (list): list of nodes that are excluded (not allowed to match as parent).
|
||||
|
||||
Returns:
|
||||
parent: The matched parent node. None if not found.
|
||||
index: The input index of matched parent node. None if not found.
|
||||
"""
|
||||
if output_name_to_node is None:
|
||||
output_name_to_node = self.model.output_name_to_node()
|
||||
|
||||
for i, inp in enumerate(node.input):
|
||||
if inp in output_name_to_node:
|
||||
parent = output_name_to_node[inp]
|
||||
if parent.op_type == parent_op_type and parent not in exclude:
|
||||
return parent, i
|
||||
|
||||
return None, None
|
||||
|
||||
def match_parent(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
parent_op_type: str,
|
||||
input_index: int | None = None,
|
||||
output_name_to_node: dict[str, onnx.NodeProto] | None = None,
|
||||
exclude: list[onnx.NodeProto] = [], # noqa: B006
|
||||
return_indice: list[int] | None = None,
|
||||
) -> onnx.NodeProto | None:
|
||||
"""
|
||||
Find parent node based on constraints on op_type and index.
|
||||
When input_index is None, we will find the first parent node based on constraints,
|
||||
and return_indice will be appended the corresponding input index.
|
||||
|
||||
Args:
|
||||
node (str): current node name.
|
||||
parent_op_type (str): constraint of parent node op_type.
|
||||
input_index (int or None): only check the parent given input index of current node.
|
||||
output_name_to_node (dict): dictionary with output name as key, and node as value.
|
||||
exclude (list): list of nodes that are excluded (not allowed to match as parent).
|
||||
return_indice (list): a list to append the input index when input_index is None.
|
||||
|
||||
Returns:
|
||||
parent: The matched parent node.
|
||||
"""
|
||||
assert node is not None
|
||||
assert input_index is None or input_index >= 0
|
||||
|
||||
if output_name_to_node is None:
|
||||
output_name_to_node = self.model.output_name_to_node()
|
||||
|
||||
if input_index is None:
|
||||
parent, index = self.match_first_parent(node, parent_op_type, output_name_to_node, exclude)
|
||||
if return_indice is not None:
|
||||
return_indice.append(index)
|
||||
return parent
|
||||
|
||||
if input_index >= len(node.input):
|
||||
# Input index out of bounds.
|
||||
return None
|
||||
|
||||
parent = self.model.get_parent(node, input_index, output_name_to_node)
|
||||
if parent is not None and parent.op_type == parent_op_type and parent not in exclude:
|
||||
return parent
|
||||
|
||||
return None
|
||||
|
||||
def match_parent_path(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
parent_op_types: list[str],
|
||||
parent_input_index: list[int] | None = None,
|
||||
output_name_to_node: dict[str, onnx.NodeProto] | None = None,
|
||||
return_indice: list[int] | None = None,
|
||||
) -> list[onnx.NodeProto] | None:
|
||||
"""
|
||||
Find a sequence of input edges based on constraints on parent op_type and index.
|
||||
When input_index is None, we will find the first parent node based on constraints,
|
||||
and return_indice will be appended the corresponding input index.
|
||||
|
||||
Args:
|
||||
node (str): current node name.
|
||||
parent_op_types (str): constraint of parent node op_type of each input edge.
|
||||
parent_input_index (list): constraint of input index of each input edge. None means no constraint.
|
||||
output_name_to_node (dict): dictionary with output name as key, and node as value.
|
||||
return_indice (list): a list to append the input index
|
||||
When there is no constraint on input index of an edge.
|
||||
|
||||
Returns:
|
||||
parents: a list of matched parent node.
|
||||
"""
|
||||
if parent_input_index is not None:
|
||||
assert len(parent_input_index) == len(parent_op_types)
|
||||
|
||||
if output_name_to_node is None:
|
||||
output_name_to_node = self.model.output_name_to_node()
|
||||
|
||||
current_node = node
|
||||
matched_parents = []
|
||||
for i, op_type in enumerate(parent_op_types):
|
||||
matched_parent = self.match_parent(
|
||||
current_node,
|
||||
op_type,
|
||||
parent_input_index[i] if parent_input_index is not None else None,
|
||||
output_name_to_node,
|
||||
exclude=[],
|
||||
return_indice=return_indice,
|
||||
)
|
||||
if matched_parent is None:
|
||||
return None
|
||||
|
||||
matched_parents.append(matched_parent)
|
||||
current_node = matched_parent
|
||||
|
||||
return matched_parents
|
||||
|
||||
def match_parent_paths(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
paths: list[tuple[list[str], list[int]]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
) -> tuple[int, list[onnx.NodeProto] | None, list[int] | None]:
|
||||
"""
|
||||
Find a matching parent path to the given node.
|
||||
"""
|
||||
for i, path in enumerate(paths):
|
||||
return_indice = []
|
||||
matched = self.match_parent_path(node, path[0], path[1], output_name_to_node, return_indice)
|
||||
if matched:
|
||||
return i, matched, return_indice
|
||||
return -1, None, None
|
||||
|
||||
def find_first_child_by_type(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
child_type: str,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]] | None = None,
|
||||
recursive: bool = True,
|
||||
) -> onnx.NodeProto | None:
|
||||
children = self.model.get_children(node, input_name_to_nodes)
|
||||
dq = deque(children)
|
||||
while len(dq) > 0:
|
||||
current_node = dq.pop()
|
||||
if current_node.op_type == child_type:
|
||||
return current_node
|
||||
|
||||
if recursive:
|
||||
children = self.model.get_children(current_node, input_name_to_nodes)
|
||||
for child in children:
|
||||
dq.appendleft(child)
|
||||
|
||||
return None
|
||||
@@ -0,0 +1,272 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
import onnx
|
||||
|
||||
from ..onnx_model import ONNXModel
|
||||
from .fusion import Fusion
|
||||
|
||||
|
||||
class FusionGelu(Fusion):
|
||||
def __init__(self, model: ONNXModel):
|
||||
super().__init__(model, "Gelu", "Erf")
|
||||
|
||||
def fuse(
|
||||
self,
|
||||
erf_node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
):
|
||||
"""
|
||||
Interface function that tries to fuse a node sequence containing an Erf node into a single
|
||||
Gelu node.
|
||||
"""
|
||||
if (
|
||||
self.fuse_1(erf_node, input_name_to_nodes, output_name_to_node)
|
||||
or self.fuse_2(erf_node, input_name_to_nodes, output_name_to_node)
|
||||
or self.fuse_3(erf_node, input_name_to_nodes, output_name_to_node)
|
||||
):
|
||||
self.model.set_opset_import("com.microsoft", 1)
|
||||
|
||||
def fuse_1(
|
||||
self,
|
||||
erf_node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
) -> bool:
|
||||
"""
|
||||
This pattern is from PyTorch model
|
||||
Fuse Gelu with Erf into one node:
|
||||
Pattern 1:
|
||||
+-------Mul(0.5)---------------------+
|
||||
| |
|
||||
| v
|
||||
[root] --> Div -----> Erf --> Add --> Mul -->
|
||||
(B=1.4142...) (1)
|
||||
|
||||
Pattern 2:
|
||||
+------------------------------------+
|
||||
| |
|
||||
| v
|
||||
[root] --> Div -----> Erf --> Add --> Mul -->Mul -->
|
||||
(B=1.4142...) (1) (0.5)
|
||||
|
||||
Note that constant input for Add and Mul could be first or second input: like either A=0.5 or B=0.5 is fine.
|
||||
"""
|
||||
if erf_node.output[0] not in input_name_to_nodes:
|
||||
return False
|
||||
children = input_name_to_nodes[erf_node.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Add":
|
||||
return False
|
||||
add_after_erf = children[0]
|
||||
|
||||
if not self.has_constant_input(add_after_erf, 1):
|
||||
return False
|
||||
|
||||
if add_after_erf.output[0] not in input_name_to_nodes:
|
||||
return False
|
||||
|
||||
children = input_name_to_nodes[add_after_erf.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Mul":
|
||||
return False
|
||||
|
||||
mul_after_erf = children[0]
|
||||
|
||||
div = self.match_parent(erf_node, "Div", 0, output_name_to_node)
|
||||
if div is None:
|
||||
return False
|
||||
|
||||
if self.find_constant_input(div, 1.4142, delta=0.001) != 1:
|
||||
return False
|
||||
|
||||
subgraph_input = div.input[0]
|
||||
|
||||
another = 1 if mul_after_erf.input[0] == add_after_erf.output[0] else 0
|
||||
if subgraph_input == mul_after_erf.input[another]: # pattern 2
|
||||
children = input_name_to_nodes[mul_after_erf.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Mul":
|
||||
return False
|
||||
mul_half = children[0]
|
||||
if not self.has_constant_input(mul_half, 0.5):
|
||||
return False
|
||||
subgraph_output = mul_half.output[0]
|
||||
else: # pattern 1
|
||||
mul_half = self.match_parent(mul_after_erf, "Mul", another, output_name_to_node)
|
||||
if mul_half is None:
|
||||
return False
|
||||
|
||||
if not self.has_constant_input(mul_half, 0.5):
|
||||
return False
|
||||
|
||||
if subgraph_input not in mul_half.input:
|
||||
return False
|
||||
|
||||
subgraph_output = mul_after_erf.output[0]
|
||||
|
||||
subgraph_nodes = [div, erf_node, add_after_erf, mul_after_erf, mul_half]
|
||||
if not self.is_safe_to_fuse_nodes(subgraph_nodes, [subgraph_output], input_name_to_nodes, output_name_to_node):
|
||||
return False
|
||||
|
||||
self.nodes_to_remove.extend(subgraph_nodes)
|
||||
fused_node = onnx.helper.make_node(
|
||||
"Gelu", name=self.create_unique_node_name(), inputs=[subgraph_input], outputs=[subgraph_output]
|
||||
)
|
||||
fused_node.domain = "com.microsoft"
|
||||
self.nodes_to_add.append(fused_node)
|
||||
return True
|
||||
|
||||
def fuse_2(
|
||||
self,
|
||||
erf_node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
) -> bool:
|
||||
"""
|
||||
This pattern is from Keras model
|
||||
Fuse Gelu with Erf into one node:
|
||||
+------------------------------------------+
|
||||
| |
|
||||
| v
|
||||
[root] --> Div -----> Erf --> Add --> Mul -->Mul
|
||||
(B=1.4142...) (A=1) (A=0.5)
|
||||
|
||||
Note that constant input for Add and Mul could be first or second input: like either A=0.5 or B=0.5 is fine.
|
||||
"""
|
||||
if erf_node.output[0] not in input_name_to_nodes:
|
||||
return False
|
||||
children = input_name_to_nodes[erf_node.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Add":
|
||||
return False
|
||||
add_after_erf = children[0]
|
||||
|
||||
if not self.has_constant_input(add_after_erf, 1):
|
||||
return False
|
||||
|
||||
if add_after_erf.output[0] not in input_name_to_nodes:
|
||||
return False
|
||||
children = input_name_to_nodes[add_after_erf.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Mul":
|
||||
return False
|
||||
mul_after_erf = children[0]
|
||||
|
||||
if not self.has_constant_input(mul_after_erf, 0.5):
|
||||
return False
|
||||
|
||||
if mul_after_erf.output[0] not in input_name_to_nodes:
|
||||
return False
|
||||
children = input_name_to_nodes[mul_after_erf.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Mul":
|
||||
return False
|
||||
mul = children[0]
|
||||
|
||||
div = self.match_parent(erf_node, "Div", 0, output_name_to_node)
|
||||
if div is None:
|
||||
return False
|
||||
|
||||
sqrt_node = None
|
||||
if self.find_constant_input(div, 1.4142, delta=0.001) != 1:
|
||||
sqrt_node = self.match_parent(div, "Sqrt", 1, output_name_to_node)
|
||||
if sqrt_node is None:
|
||||
return False
|
||||
if not self.has_constant_input(sqrt_node, 2.0):
|
||||
return False
|
||||
|
||||
subgraph_input = div.input[0]
|
||||
|
||||
if subgraph_input not in mul.input:
|
||||
return False
|
||||
|
||||
subgraph_nodes = [div, erf_node, add_after_erf, mul_after_erf, mul]
|
||||
if sqrt_node:
|
||||
subgraph_nodes.append(sqrt_node)
|
||||
|
||||
if not self.is_safe_to_fuse_nodes(subgraph_nodes, [mul.output[0]], input_name_to_nodes, output_name_to_node):
|
||||
return False
|
||||
|
||||
self.nodes_to_remove.extend(subgraph_nodes)
|
||||
fused_node = onnx.helper.make_node(
|
||||
"Gelu", name=self.create_unique_node_name(), inputs=[subgraph_input], outputs=[mul.output[0]]
|
||||
)
|
||||
fused_node.domain = "com.microsoft"
|
||||
self.nodes_to_add.append(fused_node)
|
||||
return True
|
||||
|
||||
def fuse_3(
|
||||
self,
|
||||
erf_node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
) -> bool:
|
||||
"""
|
||||
This pattern is from TensorFlow model
|
||||
Fuse Gelu with Erf into one node:
|
||||
+----------------------------------------------+
|
||||
| |
|
||||
| v
|
||||
[root] --> Mul -----> Erf --> Add --> Mul -->Mul
|
||||
(A=0.7071067690849304) (B=1) (B=0.5)
|
||||
|
||||
Note that constant input for Add and Mul could be first or second input: like either A=0.5 or B=0.5 is fine.
|
||||
"""
|
||||
|
||||
if erf_node.output[0] not in input_name_to_nodes:
|
||||
return False
|
||||
children = input_name_to_nodes[erf_node.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Add":
|
||||
return False
|
||||
add_after_erf = children[0]
|
||||
|
||||
if not self.has_constant_input(add_after_erf, 1):
|
||||
return False
|
||||
|
||||
if add_after_erf.output[0] not in input_name_to_nodes:
|
||||
return False
|
||||
children = input_name_to_nodes[add_after_erf.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Mul":
|
||||
return False
|
||||
mul_half = children[0]
|
||||
|
||||
if not self.has_constant_input(mul_half, 0.5):
|
||||
return False
|
||||
|
||||
first_mul = self.match_parent(erf_node, "Mul", 0, output_name_to_node)
|
||||
if first_mul is None:
|
||||
return False
|
||||
|
||||
i = self.find_constant_input(first_mul, 0.7071067690849304, delta=0.001)
|
||||
if i < 0:
|
||||
return False
|
||||
|
||||
root_input_index = 1 - i
|
||||
subgraph_input = first_mul.input[root_input_index]
|
||||
|
||||
if mul_half.output[0] not in input_name_to_nodes:
|
||||
return False
|
||||
children = input_name_to_nodes[mul_half.output[0]]
|
||||
if len(children) != 1 or children[0].op_type != "Mul":
|
||||
return False
|
||||
last_mul = children[0]
|
||||
|
||||
if not (last_mul.input[0] == subgraph_input or last_mul.input[1] == subgraph_input):
|
||||
return False
|
||||
|
||||
subgraph_nodes = [first_mul, erf_node, add_after_erf, mul_half, last_mul]
|
||||
if not self.is_safe_to_fuse_nodes(
|
||||
subgraph_nodes,
|
||||
[last_mul.output[0]],
|
||||
input_name_to_nodes,
|
||||
output_name_to_node,
|
||||
):
|
||||
return False
|
||||
|
||||
self.nodes_to_remove.extend(subgraph_nodes)
|
||||
fused_node = onnx.helper.make_node(
|
||||
"Gelu", name=self.create_unique_node_name(), inputs=[subgraph_input], outputs=[last_mul.output[0]]
|
||||
)
|
||||
fused_node.domain = "com.microsoft"
|
||||
self.nodes_to_add.append(fused_node)
|
||||
return True
|
||||
@@ -0,0 +1,135 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
import onnx
|
||||
|
||||
from ..onnx_model import ONNXModel
|
||||
from .fusion import Fusion
|
||||
|
||||
|
||||
class FusionLayerNormalization(Fusion):
|
||||
def __init__(self, model: ONNXModel):
|
||||
super().__init__(model, "LayerNormalization", "ReduceMean")
|
||||
|
||||
def fuse(
|
||||
self,
|
||||
reduce_mean_node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
):
|
||||
"""
|
||||
Interface function that tries to fuse a node sequence containing a ReduceMean node into a single
|
||||
LayerNormalization node.
|
||||
|
||||
+----------------------+
|
||||
| |
|
||||
| v
|
||||
[Root] --> ReduceMean --> Sub --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add
|
||||
(axis=2 or -1) | (Y=2) (axis=2 or -1) (E-6 or E-12 or 0) ^
|
||||
| |
|
||||
+-------------------------------------------------+
|
||||
|
||||
It also handles cases of duplicated sub nodes exported from older version of PyTorch:
|
||||
|
||||
+----------------------+
|
||||
| v
|
||||
| +-------> Sub-----------------------------------------------+
|
||||
| | |
|
||||
| | v
|
||||
[Root] --> ReduceMean --> Sub --> Pow --> ReduceMean --> Add --> Sqrt --> Div --> Mul --> Add
|
||||
| ^
|
||||
| |
|
||||
+----------------------+
|
||||
"""
|
||||
children = self.model.get_children(reduce_mean_node, input_name_to_nodes)
|
||||
if len(children) == 0 or len(children) > 2:
|
||||
return
|
||||
|
||||
root_input = reduce_mean_node.input[0]
|
||||
|
||||
if children[0].op_type != "Sub" or children[0].input[0] != root_input:
|
||||
return
|
||||
|
||||
if len(children) == 2:
|
||||
if children[1].op_type != "Sub" or children[1].input[0] != root_input:
|
||||
return
|
||||
|
||||
div_node = None
|
||||
for child in children:
|
||||
div_node = self.find_first_child_by_type(child, "Div", input_name_to_nodes, recursive=False)
|
||||
if div_node is not None:
|
||||
break
|
||||
if div_node is None:
|
||||
return
|
||||
|
||||
path_id, parent_nodes, _ = self.match_parent_paths(
|
||||
div_node,
|
||||
[
|
||||
(["Sqrt", "Add", "ReduceMean", "Pow", "Sub"], [1, 0, 0, 0, 0]),
|
||||
(
|
||||
["Sqrt", "Add", "ReduceMean", "Pow", "Cast", "Sub"],
|
||||
[1, 0, 0, 0, 0, 0],
|
||||
),
|
||||
],
|
||||
output_name_to_node,
|
||||
)
|
||||
if path_id < 0:
|
||||
return
|
||||
|
||||
sub_node = parent_nodes[-1]
|
||||
if sub_node not in children:
|
||||
return
|
||||
|
||||
second_add_node = parent_nodes[1]
|
||||
i, add_weight = self.get_constant_input(second_add_node)
|
||||
if add_weight is None or add_weight <= 0 or add_weight > 1.0e-4:
|
||||
# Skip fusion since epsilon value is not expected.
|
||||
return
|
||||
|
||||
pow_node = parent_nodes[3]
|
||||
if self.find_constant_input(pow_node, 2.0) != 1:
|
||||
return
|
||||
|
||||
mul_node = input_name_to_nodes[div_node.output[0]][0]
|
||||
if mul_node.op_type != "Mul":
|
||||
return
|
||||
|
||||
last_add_node = input_name_to_nodes[mul_node.output[0]][0]
|
||||
if last_add_node.op_type != "Add":
|
||||
return
|
||||
|
||||
subgraph_nodes = [reduce_mean_node]
|
||||
subgraph_nodes.extend(children)
|
||||
subgraph_nodes.extend(parent_nodes[:-1])
|
||||
|
||||
subgraph_nodes.extend([last_add_node, mul_node, div_node])
|
||||
if not self.is_safe_to_fuse_nodes(
|
||||
subgraph_nodes,
|
||||
last_add_node.output,
|
||||
input_name_to_nodes,
|
||||
output_name_to_node,
|
||||
):
|
||||
return
|
||||
|
||||
weight_input = mul_node.input[1 - self.input_index(div_node.output[0], mul_node)]
|
||||
if not self.is_constant_with_specified_rank(weight_input, 1):
|
||||
return
|
||||
|
||||
bias_input = last_add_node.input[1 - self.input_index(mul_node.output[0], last_add_node)]
|
||||
if not self.is_constant_with_specified_rank(bias_input, 1):
|
||||
return
|
||||
|
||||
self.nodes_to_remove.extend(subgraph_nodes)
|
||||
|
||||
normalize_node = onnx.helper.make_node(
|
||||
"LayerNormalization",
|
||||
name=self.create_unique_node_name(),
|
||||
inputs=[reduce_mean_node.input[0], weight_input, bias_input],
|
||||
outputs=[last_add_node.output[0]],
|
||||
)
|
||||
normalize_node.attribute.extend([onnx.helper.make_attribute("epsilon", float(add_weight))])
|
||||
self.nodes_to_add.append(normalize_node)
|
||||
@@ -0,0 +1,96 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
|
||||
from ..onnx_model import ONNXModel
|
||||
from .fusion import Fusion
|
||||
|
||||
|
||||
class ReplaceUpsampleWithResize(Fusion):
|
||||
"""Replace Upsample with Resize."""
|
||||
|
||||
def __init__(self, model: ONNXModel, opset):
|
||||
"""Initialize."""
|
||||
super().__init__(model, "Resize", "Upsample")
|
||||
self.opset = opset
|
||||
|
||||
def fuse(
|
||||
self,
|
||||
node: onnx.NodeProto,
|
||||
input_name_to_nodes: dict[str, list[onnx.NodeProto]],
|
||||
output_name_to_node: dict[str, onnx.NodeProto],
|
||||
):
|
||||
"""Replace Upsample with Resize."""
|
||||
mode = None
|
||||
for attr in node.attribute:
|
||||
if attr.name == "mode":
|
||||
mode = attr.s.decode("utf-8")
|
||||
break
|
||||
|
||||
scales_input = None
|
||||
if self.opset > 7:
|
||||
scales_input = node.input[1] if len(node.input) > 1 else ""
|
||||
resize_inputs = [node.input[0], node.name + "_roi", scales_input]
|
||||
else:
|
||||
if self.opset == 7:
|
||||
for attr in node.attribute:
|
||||
if attr.name == "scales":
|
||||
scales_input = attr.floats
|
||||
break
|
||||
|
||||
scales_input = np.array(list(scales_input), np.float32)
|
||||
else:
|
||||
h_scale = 1
|
||||
w_scale = 1
|
||||
for attr in node.attribute:
|
||||
if attr.name == "height_scale":
|
||||
h_scale = attr.float
|
||||
elif attr.name == "width_scale":
|
||||
w_scale = attr.float
|
||||
|
||||
scales_input = np.array([1, 1, h_scale, w_scale], np.float32)
|
||||
|
||||
scales_tensor = onnx.helper.make_tensor(
|
||||
name=node.name + "_scales",
|
||||
data_type=onnx.TensorProto.FLOAT,
|
||||
dims=scales_input.shape,
|
||||
vals=scales_input.flatten().tolist(),
|
||||
)
|
||||
|
||||
scales_node = onnx.helper.make_node(
|
||||
"Constant", inputs=[], outputs=[node.name + "_scales"], value=scales_tensor
|
||||
)
|
||||
|
||||
self.nodes_to_add.append(scales_node)
|
||||
|
||||
resize_inputs = [node.input[0], node.name + "_roi", node.name + "_scales"]
|
||||
|
||||
roi_tensor = onnx.helper.make_tensor(
|
||||
name=node.name + "_roi",
|
||||
data_type=onnx.TensorProto.FLOAT,
|
||||
dims=(len(scales_input) * 2,),
|
||||
vals=[0] * len(scales_input) + [1] * len(scales_input),
|
||||
)
|
||||
|
||||
roi_node = onnx.helper.make_node("Constant", inputs=[], outputs=[node.name + "_roi"], value=roi_tensor)
|
||||
|
||||
resize_node = onnx.helper.make_node(
|
||||
op_type="Resize", inputs=resize_inputs, outputs=node.output, mode=mode, nearest_mode="floor"
|
||||
)
|
||||
|
||||
self.nodes_to_remove.append(node)
|
||||
self.nodes_to_add.append(roi_node)
|
||||
self.nodes_to_add.append(resize_node)
|
||||
|
||||
def apply(self) -> bool:
|
||||
"""Apply."""
|
||||
if super().apply():
|
||||
self.model.topological_sort()
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,239 @@
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import onnx
|
||||
from onnx.onnx_pb import GraphProto, ModelProto, NodeProto, TensorProto
|
||||
|
||||
from onnxruntime.capi._pybind_state import quantize_matmul_bnb4
|
||||
|
||||
from .onnx_model import ONNXModel
|
||||
from .quant_utils import attribute_to_kwarg
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MatMulBnb4Quantizer:
|
||||
"""Perform 4b quantization of constant MatMul weights using FP4 or NF4 data type"""
|
||||
|
||||
##################
|
||||
# quantization types, must be consistent with native code type
|
||||
# Bnb_DataType_t defined in blockwise_quant_block_bnb4.h
|
||||
|
||||
# 4b floating point with bias of 3
|
||||
FP4 = 0
|
||||
|
||||
# 4b NormalFloat
|
||||
NF4 = 1
|
||||
|
||||
def __init__(self, model: ModelProto, quant_type: int, block_size: int, nodes_to_exclude=None):
|
||||
nodes_to_exclude = nodes_to_exclude or []
|
||||
assert quant_type in [MatMulBnb4Quantizer.FP4, MatMulBnb4Quantizer.NF4]
|
||||
self.model = ONNXModel(model)
|
||||
self.quant_type = quant_type
|
||||
self.block_size = block_size
|
||||
self.nodes_to_exclude = set(nodes_to_exclude)
|
||||
|
||||
@staticmethod
|
||||
def __get_initializer(name, graph_path: list[GraphProto]) -> tuple[TensorProto, GraphProto]:
|
||||
for gid in range(len(graph_path) - 1, -1, -1):
|
||||
graph = graph_path[gid]
|
||||
for tensor in graph.initializer:
|
||||
if tensor.name == name:
|
||||
return tensor, graph
|
||||
return None, None
|
||||
|
||||
def bnb4_block_quant(self, fpweight: npt.ArrayLike) -> np.ndarray:
|
||||
"""4b quantize fp32/fp16 weight"""
|
||||
|
||||
if len(fpweight.shape) != 2:
|
||||
raise ValueError("Current bnb4 block quantization only supports 2D tensors!")
|
||||
# need to copy since the transposed weight still has the original memory layout
|
||||
# Linear4bit quantizes its weight data which is the transposed weight
|
||||
fpweight_t = fpweight.transpose().copy()
|
||||
|
||||
rows, cols = fpweight.shape
|
||||
numel = rows * cols
|
||||
block_size = self.block_size
|
||||
num_blocks = (numel + block_size - 1) // block_size
|
||||
quantized_numel = (numel + 1) // 2
|
||||
|
||||
packed = np.zeros(quantized_numel, dtype="uint8")
|
||||
absmax = np.zeros(num_blocks, dtype=fpweight.dtype)
|
||||
# block wise quantization, fpweight_t is flattened and divided into blocks
|
||||
quantize_matmul_bnb4(packed, fpweight_t, absmax, block_size, self.quant_type, cols, rows)
|
||||
|
||||
return (packed, absmax)
|
||||
|
||||
def _bnb4_matmul_node_weight(self, node: NodeProto, graph_stack: list[GraphProto]) -> NodeProto:
|
||||
"""If the node is MatMul with fp32 const weight, quantize the weight with int4, and return the new node"""
|
||||
|
||||
if node.op_type != "MatMul":
|
||||
return node # only care about MatMul for now
|
||||
|
||||
logger.debug(f"start to quantize {node.name} ...")
|
||||
if node.name in self.nodes_to_exclude:
|
||||
logger.debug(f"exclude to quantize {node.name} as specified by nodes_to_exclude...")
|
||||
return node
|
||||
|
||||
inputB = node.input[1] # noqa: N806
|
||||
B, Bs_graph = MatMulBnb4Quantizer.__get_initializer(inputB, graph_stack) # noqa: N806
|
||||
if B is None:
|
||||
logger.debug("MatMul doesn't have const weight. Skip to quantize")
|
||||
return node # only care about constant weight
|
||||
|
||||
B_array = onnx.numpy_helper.to_array(B) # noqa: N806
|
||||
if len(B_array.shape) != 2:
|
||||
logger.debug("MatMul weight is not 2D. Skip to quantize")
|
||||
return node # can only process 2-D matrix
|
||||
|
||||
packed, absmax = self.bnb4_block_quant(B_array)
|
||||
B_quant = onnx.numpy_helper.from_array(packed) # noqa: N806
|
||||
B_quant.name = B.name + "_Bnb4"
|
||||
for input in Bs_graph.input:
|
||||
if input.name == inputB:
|
||||
Bs_graph.input.remove(input)
|
||||
break
|
||||
|
||||
absmax_tensor = onnx.numpy_helper.from_array(absmax)
|
||||
absmax_tensor.name = B.name + "_absmax"
|
||||
|
||||
Bs_graph.initializer.extend([B_quant, absmax_tensor])
|
||||
|
||||
kwargs = {}
|
||||
rows, cols = B_array.shape
|
||||
kwargs["K"] = rows
|
||||
kwargs["N"] = cols
|
||||
kwargs["block_size"] = self.block_size
|
||||
kwargs["quant_type"] = self.quant_type
|
||||
|
||||
matmul_bnb4_node = onnx.helper.make_node(
|
||||
"MatMulBnb4",
|
||||
inputs=[node.input[0], B_quant.name, absmax_tensor.name],
|
||||
outputs=[node.output[0]],
|
||||
name=node.name + "_Bnb4" if node.name else "",
|
||||
domain="com.microsoft",
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
logger.debug(f"complete quantization of {node.name} ...")
|
||||
|
||||
return matmul_bnb4_node
|
||||
|
||||
def _process_subgraph(self, graph_stack: list[GraphProto]):
|
||||
new_nodes = []
|
||||
graph = graph_stack[-1]
|
||||
|
||||
for node in graph.node:
|
||||
graph_attrs = [
|
||||
attr
|
||||
for attr in node.attribute
|
||||
if attr.type == onnx.AttributeProto.GRAPH or attr.type == onnx.AttributeProto.GRAPHS
|
||||
]
|
||||
if graph_attrs:
|
||||
kwargs = {}
|
||||
for attr in node.attribute:
|
||||
if attr.type == onnx.AttributeProto.GRAPH:
|
||||
# recursive call to take care of sub-graph
|
||||
graph_stack.append(attr.g)
|
||||
kv = {attr.name: self._process_subgraph(graph_stack)}
|
||||
elif attr.type == onnx.AttributeProto.GRAPHS:
|
||||
value = []
|
||||
for subgraph in attr.graphs:
|
||||
# recursive call to take care of sub-graph
|
||||
graph_stack.append(subgraph)
|
||||
value.extend([self._process_subgraph(graph_stack)])
|
||||
kv = {attr.name: value}
|
||||
else:
|
||||
kv = attribute_to_kwarg(attr)
|
||||
kwargs.update(kv)
|
||||
node = onnx.helper.make_node( # noqa: PLW2901
|
||||
node.op_type, node.input, node.output, name=node.name, **kwargs
|
||||
)
|
||||
|
||||
new_nodes.append(self._bnb4_matmul_node_weight(node, graph_stack))
|
||||
|
||||
graph.ClearField("node")
|
||||
graph.node.extend(new_nodes)
|
||||
graph_stack.pop()
|
||||
return graph
|
||||
|
||||
def process(self):
|
||||
# use a stack to keep track of sub-graphs
|
||||
graph_stack = [self.model.graph()]
|
||||
opset_import = self.model.opset_import()
|
||||
|
||||
has_ms_domain = False
|
||||
for opset in opset_import:
|
||||
if opset.domain == "com.microsoft":
|
||||
has_ms_domain = True
|
||||
if not has_ms_domain:
|
||||
opset_import.extend([onnx.helper.make_opsetid("com.microsoft", 1)])
|
||||
|
||||
self._process_subgraph(graph_stack)
|
||||
self.model.clean_initializers()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="""Blockwise FP4/NF4 quantization for MatMul 2D weight matrices.
|
||||
|
||||
A weight matrix is partitioned into blocks, where each block is a contiguous
|
||||
subset inside the flattened transposed weight matrix. Each block is quantized
|
||||
into a set of 4b integers with an absolute value scaling factor.
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument("--input_model", required=True, help="Path to the input model file")
|
||||
parser.add_argument("--output_model", required=True, help="Path to the output model file")
|
||||
parser.add_argument(
|
||||
"--quant_type",
|
||||
required=False,
|
||||
default=1,
|
||||
choices=[MatMulBnb4Quantizer.FP4, MatMulBnb4Quantizer.NF4],
|
||||
help="Quantization data type. 0: FP4, 1: NF4",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--block_size",
|
||||
required=False,
|
||||
default=64,
|
||||
help="Block size for blockwise quantization. Note: bnb.nn.Linear4bit only uses block_size=64",
|
||||
)
|
||||
parser.add_argument("-v", "--verbose", required=False, action="store_true")
|
||||
parser.set_defaults(verbose=False)
|
||||
parser.add_argument(
|
||||
"--nodes_to_exclude",
|
||||
nargs="+",
|
||||
type=str,
|
||||
required=False,
|
||||
default=[],
|
||||
help="Specify the nodes to be excluded from quantization with node names",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
if args.verbose:
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
input_model_path = args.input_model
|
||||
output_model_path = args.output_model
|
||||
|
||||
if os.path.exists(output_model_path):
|
||||
logger.error(f"file {output_model_path} already exists")
|
||||
raise Exception(f"file {output_model_path} already exists")
|
||||
|
||||
model = onnx.load(input_model_path)
|
||||
quant = MatMulBnb4Quantizer(model, args.quant_type, args.block_size, nodes_to_exclude=args.nodes_to_exclude)
|
||||
quant.process()
|
||||
quant.model.save_model_to_file(output_model_path, True)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
from .weight_only import gptq_quantize, rtn_quantize # noqa: F401
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,80 @@
|
||||
#
|
||||
# The implementation of this file is based on:
|
||||
# https://github.com/intel/neural-compressor/tree/master/neural_compressor
|
||||
#
|
||||
# Copyright (c) 2023 Intel Corporation
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Helper classes or functions for onnxrt adaptor."""
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger("neural_compressor")
|
||||
|
||||
|
||||
MAXIMUM_PROTOBUF = 2147483648
|
||||
|
||||
|
||||
def simple_progress_bar(total, i):
|
||||
"""Progress bar for cases where tqdm can't be used."""
|
||||
progress = i / total
|
||||
bar_length = 20
|
||||
bar = "#" * int(bar_length * progress)
|
||||
spaces = " " * (bar_length - len(bar))
|
||||
percentage = progress * 100
|
||||
print(f"\rProgress: [{bar}{spaces}] {percentage:.2f}%", end="")
|
||||
|
||||
|
||||
def find_by_name(name, item_list):
|
||||
"""Helper function to find item by name in a list."""
|
||||
items = []
|
||||
for item in item_list:
|
||||
assert hasattr(item, "name"), f"{item} should have a 'name' attribute defined" # pragma: no cover
|
||||
if item.name == name:
|
||||
items.append(item)
|
||||
if len(items) > 0:
|
||||
return items[0]
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def to_numpy(data):
|
||||
"""Convert to numpy ndarrays."""
|
||||
import torch # noqa: PLC0415
|
||||
|
||||
if not isinstance(data, np.ndarray):
|
||||
if not importlib.util.find_spec("torch"):
|
||||
logger.error(
|
||||
"Please install torch to enable subsequent data type check and conversion, "
|
||||
"or reorganize your data format to numpy array."
|
||||
)
|
||||
exit(0)
|
||||
if isinstance(data, torch.Tensor):
|
||||
if data.dtype is torch.bfloat16: # pragma: no cover
|
||||
return data.detach().cpu().to(torch.float32).numpy()
|
||||
if data.dtype is torch.chalf: # pragma: no cover
|
||||
return data.detach().cpu().to(torch.cfloat).numpy()
|
||||
return data.detach().cpu().numpy()
|
||||
else:
|
||||
try:
|
||||
return np.array(data)
|
||||
except Exception:
|
||||
assert False, ( # noqa: B011
|
||||
f"The input data for onnx model is {type(data)}, which is not supported to convert to numpy ndarrays."
|
||||
)
|
||||
else:
|
||||
return data
|
||||
@@ -0,0 +1,932 @@
|
||||
#
|
||||
# The implementation of this file is based on:
|
||||
# https://github.com/intel/neural-compressor/tree/master/neural_compressor
|
||||
#
|
||||
# Copyright (c) 2023 Intel Corporation
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
# Modifications:
|
||||
# Add k-quant quantization method.
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
|
||||
"""WeightOnly for onnxrt adaptor."""
|
||||
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
from onnx import numpy_helper
|
||||
from onnx.helper import np_dtype_to_tensor_dtype
|
||||
|
||||
import onnxruntime as ort
|
||||
|
||||
from .onnx_model import ONNXModel
|
||||
from .util import simple_progress_bar
|
||||
|
||||
logger = logging.getLogger("neural_compressor")
|
||||
|
||||
|
||||
def make_matmul_weight_only_node(
|
||||
node,
|
||||
weight_shape,
|
||||
num_bits,
|
||||
group_size,
|
||||
k_blocks,
|
||||
q_weight,
|
||||
scale,
|
||||
zero_point,
|
||||
accuracy_level=0,
|
||||
): # pragma: no cover
|
||||
"""Build MatMulNBits node.
|
||||
|
||||
Args:
|
||||
node: original matmul node
|
||||
weight_shape: original weight shape
|
||||
num_bits (int): num_bits
|
||||
group_size (int): how many elements share one scale/zp
|
||||
k_blocks (int): block number
|
||||
q_weight (array): quantized weight
|
||||
scale (array): scale
|
||||
zero_point (array): zero point
|
||||
accuracy_level (int): accuracy level. Support 0 (unset), 1(fp32), 2(fp16), 3(bf16), or 4(int8).
|
||||
|
||||
Returns:
|
||||
matmul_weight_only_node: MatMulNBits node
|
||||
new_inits: initializers of the new node
|
||||
"""
|
||||
blob_size = group_size * num_bits // 8
|
||||
packed = np.zeros((q_weight.shape[0], blob_size), dtype="uint8")
|
||||
q_weight_name = node.input[1] + f"_Q{num_bits!s}G{group_size!s}"
|
||||
input_names = [node.input[0], q_weight_name]
|
||||
new_inits = []
|
||||
kwargs = {}
|
||||
|
||||
op_type = "MatMulNBits"
|
||||
|
||||
# pack quantized weight
|
||||
if num_bits == 4:
|
||||
q_weight_pairs = q_weight[:, ::2] | q_weight[:, 1::2] << 4
|
||||
packed[:, :] = q_weight_pairs[:, :blob_size]
|
||||
elif num_bits == 8:
|
||||
packed = q_weight
|
||||
else:
|
||||
logger.error(f"MatMulNBits does not have kernel support for num_bits = {num_bits}.")
|
||||
|
||||
packed = np.reshape(packed, (-1, k_blocks, blob_size))
|
||||
|
||||
# build scale tensor
|
||||
scale = np.reshape(scale, (-1, k_blocks))
|
||||
assert scale.dtype == np.float32 or scale.dtype == np.float16
|
||||
scale_tensor = onnx.helper.make_tensor(
|
||||
name=node.input[1] + "_scale",
|
||||
data_type=np_dtype_to_tensor_dtype(scale.dtype),
|
||||
dims=scale.shape,
|
||||
vals=scale.tobytes(),
|
||||
raw=True,
|
||||
)
|
||||
input_names.append(scale_tensor.name)
|
||||
new_inits.append(scale_tensor)
|
||||
|
||||
# build zero_point tensor
|
||||
if zero_point is not None:
|
||||
if num_bits == 8:
|
||||
packed_zp = zero_point.astype("uint8")
|
||||
elif num_bits == 4:
|
||||
# For 4-bit case, the default zeros is 0x8. So it is 0x88 = 136 if we fill lower/higher 4 bits with 0x8.
|
||||
packed_zp = np.full((zero_point.shape[0] + 1) // 2, 136, dtype="uint8")
|
||||
# create an index array
|
||||
idx = np.arange(zero_point.shape[0] // k_blocks * k_blocks).reshape(-1)
|
||||
# separate odd and even indices
|
||||
even_idx = idx[::2]
|
||||
odd_idx = idx[1::2]
|
||||
# vectorized operation for even and odd indices
|
||||
packed_zp[even_idx // 2] = (packed_zp[even_idx // 2] & 0xF0) | zero_point[even_idx].ravel()
|
||||
packed_zp[odd_idx // 2] = (packed_zp[odd_idx // 2] & 0x0F) | (zero_point[odd_idx].ravel() << 4)
|
||||
else:
|
||||
raise ValueError(f"MatMulNBits does not have kernel support for num_bits = {num_bits}.")
|
||||
|
||||
packed_zp = np.reshape(packed_zp, (weight_shape[1], -1))
|
||||
zp_tensor = onnx.helper.make_tensor(
|
||||
name=node.input[1] + "_zp", data_type=2, dims=packed_zp.shape, vals=packed_zp.tobytes(), raw=True
|
||||
)
|
||||
input_names.append(zp_tensor.name)
|
||||
new_inits.append(zp_tensor)
|
||||
|
||||
# set kwargs
|
||||
kwargs["K"] = weight_shape[0]
|
||||
kwargs["N"] = weight_shape[1]
|
||||
kwargs["bits"] = num_bits
|
||||
kwargs["block_size"] = group_size
|
||||
if accuracy_level > 0:
|
||||
# require onnxruntime > 1.16.3
|
||||
kwargs["accuracy_level"] = accuracy_level
|
||||
|
||||
q_weight_tensor = onnx.helper.make_tensor(
|
||||
name=q_weight_name,
|
||||
data_type=2,
|
||||
dims=packed.shape,
|
||||
vals=packed.tobytes(),
|
||||
raw=True,
|
||||
)
|
||||
new_inits.append(q_weight_tensor)
|
||||
|
||||
matmul_weight_only_node = onnx.helper.make_node(
|
||||
op_type,
|
||||
inputs=input_names,
|
||||
outputs=node.output,
|
||||
name=node.name + "_Q" + str(num_bits) if node.name else "_Q" + str(num_bits),
|
||||
domain="com.microsoft",
|
||||
**kwargs,
|
||||
)
|
||||
return matmul_weight_only_node, new_inits
|
||||
|
||||
|
||||
def quant_tensor(data, num_bits=4, group_size=32, scheme="asym", dtype="int", ratio=1.0):
|
||||
"""Quantize tensor per group.
|
||||
|
||||
Args:
|
||||
data : input weight
|
||||
num_bits (int, optional): num_bits. Defaults to 4.
|
||||
group_size (int, optional): how many elements share one scale/zp. Defaults to 4.
|
||||
scheme (str, optional): quantization scheme. Defaults to "asym".
|
||||
dtype (str, optional): data type. Defaults to "int".
|
||||
ratio (float, optional): percentile of clip. Defaults to 1.0.
|
||||
|
||||
Returns:
|
||||
output: quantized weight
|
||||
scale: scale
|
||||
zero_point: zero point
|
||||
"""
|
||||
data = np.reshape(data, (-1, group_size))
|
||||
if scheme == "asym" or dtype == "uint":
|
||||
maxq = 2**num_bits - 1
|
||||
minq = 0
|
||||
elif scheme == "sym":
|
||||
maxq = 2 ** (num_bits - 1) - 1 if num_bits != 1 else 0
|
||||
minq = -(2 ** (num_bits - 1)) if num_bits != 1 else -1
|
||||
|
||||
rmin = np.min(data, axis=1, keepdims=True) * ratio
|
||||
rmax = np.max(data, axis=1, keepdims=True) * ratio
|
||||
if scheme == "sym":
|
||||
max_range = np.maximum(np.abs(rmin), np.abs(rmax))
|
||||
scale = np.ones(rmax.shape)
|
||||
mask = max_range > 0
|
||||
scale[mask] = (max_range[mask] * 2.0).astype(np.float64) / (maxq - minq)
|
||||
zero_point = (
|
||||
np.zeros(scale.shape) if dtype == "int" else np.ones(rmax.shape, dtype="uint8") * (1 << (num_bits - 1))
|
||||
)
|
||||
else:
|
||||
scale = np.ones(rmax.shape)
|
||||
scale[rmin != rmax] = np.array(
|
||||
[float(i) / (maxq - minq) for i in (rmax - rmin)[rmin != rmax].flatten().tolist()]
|
||||
)
|
||||
zero_point = (
|
||||
((np.zeros(scale.shape) - rmin) / scale).round()
|
||||
if dtype == "int"
|
||||
else np.maximum(0, np.minimum(maxq, ((np.zeros(scale.shape) - rmin) / scale).round())).astype("uint8")
|
||||
)
|
||||
|
||||
q_weight = np.empty_like(data, dtype=scale.dtype)
|
||||
np.divide(data, scale, out=q_weight)
|
||||
np.add(q_weight, zero_point, out=q_weight)
|
||||
np.round(q_weight, out=q_weight)
|
||||
np.clip(q_weight, minq, maxq, out=q_weight)
|
||||
|
||||
return q_weight, scale, zero_point
|
||||
|
||||
|
||||
def quant_tensor_k_quant_cpu(data, num_bits=4, group_size=32):
|
||||
"""Quantize tensor per group based on k quant.
|
||||
|
||||
Ref: https://github.com/ggml-org/llama.cpp/blob/64eda5deb9859e87a020e56bab5d2f9ca956f1de/ggml/src/ggml-quants.c
|
||||
|
||||
Args:
|
||||
data : input weight
|
||||
num_bits (int, optional): num_bits. Defaults to 4.
|
||||
group_size (int, optional): how many elements share one scale/zp. Defaults to 32.
|
||||
|
||||
Returns:
|
||||
output: quantized weight
|
||||
scale: scale
|
||||
zero_point: zero point
|
||||
"""
|
||||
data = np.reshape(data, (-1, group_size)).astype(np.float32) # nb = data.shape[0], (nb, group_size)
|
||||
maxq = 2**num_bits - 1
|
||||
minq = 0
|
||||
sum_x2 = np.sum(data**2, axis=1, keepdims=True) # (nb, 1)
|
||||
av_x = np.sqrt(sum_x2 / group_size) # (nb, 1)
|
||||
weights = np.add(av_x, np.abs(data)) # (nb, group_size)
|
||||
rmin = np.min(data, axis=1, keepdims=True) # (nb, 1)
|
||||
rmax = np.max(data, axis=1, keepdims=True) # (nb, 1)
|
||||
sum_w = np.sum(weights, axis=1, keepdims=True) # (nb, 1)
|
||||
sum_x = np.sum(weights * data, axis=1, keepdims=True) # (nb, group_size)
|
||||
iscale = np.ones(rmax.shape, dtype=data.dtype) # (nb, 1)
|
||||
mask = rmin != rmax
|
||||
iscale[mask] = (maxq - minq) / (rmax[mask] - rmin[mask])
|
||||
scale = 1 / iscale
|
||||
quant_data = np.clip(np.round(iscale * (data - rmin)), minq, maxq) # (nb, group_size)
|
||||
diff = scale * quant_data + rmin - data # (nb, group_size)
|
||||
best_mad = np.sum(weights * diff**2, axis=1, keepdims=True) # (nb, 1)
|
||||
nstep = 20
|
||||
rdelta = 0.1
|
||||
# nstep * rdelta = -2 * rrmin, maxq - minq = 2**num_bits - 1
|
||||
rrmin = -1
|
||||
for is_ in range(nstep):
|
||||
iscale_new = np.ones(rmax.shape, dtype=data.dtype) # (nb, 1)
|
||||
factor = np.array([rrmin + rdelta * is_ + maxq - minq]).astype(data.dtype)[0]
|
||||
mask = rmin != rmax
|
||||
iscale_new[mask] = factor / (rmax[mask] - rmin[mask])
|
||||
quant_data_new = np.clip(np.round(iscale_new * (data - rmin)), minq, maxq) # (nb, group_size)
|
||||
mul_weights_quant_data_new = weights * quant_data_new
|
||||
sum_l = np.sum(mul_weights_quant_data_new, axis=1, keepdims=True) # (nb, 1)
|
||||
sum_l2 = np.sum(mul_weights_quant_data_new * quant_data_new, axis=1, keepdims=True) # (nb, 1)
|
||||
sum_xl = np.sum(mul_weights_quant_data_new * data, axis=1, keepdims=True) # (nb, 1)
|
||||
D = np.subtract(sum_w * sum_l2, sum_l**2) # noqa: N806
|
||||
|
||||
this_scale = (sum_w * sum_xl - sum_x * sum_l) / D # (nb, 1)
|
||||
this_min = (sum_l2 * sum_x - sum_l * sum_xl) / D # (nb, 1)
|
||||
|
||||
diff = this_scale * quant_data_new + this_min - data # (nb, group_size)
|
||||
mad = np.sum(weights * diff**2, axis=1, keepdims=True) # (nb, 1)
|
||||
|
||||
mad_1 = np.array(mad)
|
||||
best_mad_1 = np.array(best_mad)
|
||||
idx_to_replace = np.where(mad_1 < best_mad_1)[0]
|
||||
quant_data[idx_to_replace, :] = quant_data_new[idx_to_replace, :]
|
||||
best_mad[idx_to_replace] = mad[idx_to_replace]
|
||||
scale[idx_to_replace] = this_scale[idx_to_replace]
|
||||
rmin[idx_to_replace] = this_min[idx_to_replace]
|
||||
|
||||
zero_point = np.clip(((-rmin) / scale).round(), 0, maxq).astype("uint8")
|
||||
scale = scale.astype(np.float64)
|
||||
q_weight = np.empty_like(data, dtype=scale.dtype)
|
||||
np.divide(data, scale, out=q_weight)
|
||||
np.add(q_weight, zero_point, out=q_weight)
|
||||
np.round(q_weight, out=q_weight)
|
||||
np.clip(q_weight, minq, maxq, out=q_weight)
|
||||
|
||||
return q_weight, scale, zero_point
|
||||
|
||||
|
||||
def quant_tensor_k_quant_cuda(data, num_bits=4, group_size=32):
|
||||
"""Quantize tensor per group based on k quant.
|
||||
|
||||
Ref: https://github.com/ggml-org/llama.cpp/blob/64eda5deb9859e87a020e56bab5d2f9ca956f1de/ggml/src/ggml-quants.c
|
||||
|
||||
Args:
|
||||
data : input weight
|
||||
num_bits (int, optional): num_bits. Defaults to 4.
|
||||
group_size (int, optional): how many elements share one scale/zp. Defaults to 4.
|
||||
|
||||
Returns:
|
||||
output: quantized weight
|
||||
scale: scale
|
||||
zero_point: zero point
|
||||
"""
|
||||
try:
|
||||
import cupy as cp # noqa: PLC0415
|
||||
import torch # noqa: PLC0415
|
||||
|
||||
if torch.cuda.is_available():
|
||||
data = cp.asarray(data)
|
||||
data = data.reshape((-1, group_size)).astype(cp.float32) # nb = data.shape[0], (nb, group_size)
|
||||
maxq = 2**num_bits - 1
|
||||
minq = 0
|
||||
sum_x2 = cp.sum(data**2, axis=1, keepdims=True) # (nb, 1)
|
||||
av_x = cp.sqrt(sum_x2 / group_size) # (nb, 1)
|
||||
weights = cp.add(av_x, cp.abs(data)) # (nb, group_size)
|
||||
rmin = cp.min(data, axis=1, keepdims=True) # (nb, 1)
|
||||
rmax = cp.max(data, axis=1, keepdims=True) # (nb, 1)
|
||||
sum_w = cp.sum(weights, axis=1, keepdims=True) # (nb, 1)
|
||||
sum_x = cp.sum(weights * data, axis=1, keepdims=True) # (nb, group_size)
|
||||
iscale = cp.ones(rmax.shape, dtype=data.dtype) # (nb, 1)
|
||||
mask = rmin != rmax
|
||||
iscale[mask] = (maxq - minq) / (rmax[mask] - rmin[mask])
|
||||
scale = 1 / iscale
|
||||
quant_data = cp.clip(cp.round(iscale * (data - rmin)), minq, maxq) # (nb, group_size)
|
||||
diff = scale * quant_data + rmin - data # (nb, group_size)
|
||||
best_mad = cp.sum(weights * diff**2, axis=1, keepdims=True) # (nb, 1)
|
||||
nstep = 20
|
||||
rdelta = 0.1
|
||||
rrmin = -1
|
||||
for is_ in range(nstep):
|
||||
iscale_new = cp.ones(rmax.shape, dtype=data.dtype) # (nb, 1)
|
||||
factor = cp.array([rrmin + rdelta * is_ + maxq - minq]).astype(data.dtype)[0]
|
||||
mask = rmin != rmax
|
||||
iscale_new[mask] = factor / (rmax[mask] - rmin[mask])
|
||||
quant_data_new = cp.clip(cp.round(iscale_new * (data - rmin)), minq, maxq) # (nb, group_size)
|
||||
mul_weights_quant_data_new = weights * quant_data_new
|
||||
sum_l = cp.sum(mul_weights_quant_data_new, axis=1, keepdims=True) # (nb, 1)
|
||||
sum_l2 = cp.sum(mul_weights_quant_data_new * quant_data_new, axis=1, keepdims=True) # (nb, 1)
|
||||
sum_xl = cp.sum(mul_weights_quant_data_new * data, axis=1, keepdims=True) # (nb, 1)
|
||||
D = cp.subtract(sum_w * sum_l2, sum_l**2) # noqa: N806
|
||||
|
||||
this_scale = (sum_w * sum_xl - sum_x * sum_l) / D # (nb, 1)
|
||||
this_min = (sum_l2 * sum_x - sum_l * sum_xl) / D # (nb, 1)
|
||||
|
||||
diff = this_scale * quant_data_new + this_min - data # (nb, group_size)
|
||||
mad = cp.sum(weights * diff**2, axis=1, keepdims=True) # (nb, 1)
|
||||
|
||||
mad_1 = cp.array(mad)
|
||||
best_mad_1 = cp.array(best_mad)
|
||||
idx_to_replace = cp.where(mad_1 < best_mad_1)[0]
|
||||
quant_data[idx_to_replace, :] = quant_data_new[idx_to_replace, :]
|
||||
best_mad[idx_to_replace] = mad[idx_to_replace]
|
||||
scale[idx_to_replace] = this_scale[idx_to_replace]
|
||||
rmin[idx_to_replace] = this_min[idx_to_replace]
|
||||
|
||||
zero_point = cp.clip(((-rmin) / scale).round(), 0, maxq).astype("uint8")
|
||||
scale = scale.astype(cp.float64)
|
||||
q_weight = cp.empty_like(data, dtype=scale.dtype)
|
||||
cp.divide(data, scale, out=q_weight)
|
||||
cp.add(q_weight, zero_point, out=q_weight)
|
||||
cp.round(q_weight, out=q_weight)
|
||||
cp.clip(q_weight, minq, maxq, out=q_weight)
|
||||
|
||||
return q_weight.get(), scale.get(), zero_point.get()
|
||||
else:
|
||||
logger.warning(
|
||||
"Try to use k-quant quantization on CUDA. However, CUDA is not available."
|
||||
"Fall back to k-quant quantization on CPU."
|
||||
)
|
||||
return quant_tensor_k_quant_cpu(data, num_bits, group_size)
|
||||
except ImportError:
|
||||
logger.info(
|
||||
"Now we are using k-quant quantization on cpu, which is time consuming."
|
||||
"Please consider install cupy to speed up on CUDA. See https://cupy.dev/"
|
||||
"Please also install torch to check CUDA availability."
|
||||
)
|
||||
return quant_tensor_k_quant_cpu(data, num_bits, group_size)
|
||||
|
||||
|
||||
def qdq_tensor(data, num_bits=4, group_size=32, scheme="asym", dtype="int", ratio=1.0):
|
||||
"""Quant dequant tensor per group.
|
||||
|
||||
Args:
|
||||
data : input weight
|
||||
num_bits (int, optional): num_bits. Defaults to 4.
|
||||
group_size (int, optional): how many elements share one scale/zp. Defaults to 4.
|
||||
scheme (str, optional): quantization scheme. Defaults to "asym".
|
||||
dtype (str, optional): data type. Defaults to "int".
|
||||
ratio (float, optional): percentile of clip. Defaults to 1.0.
|
||||
|
||||
Returns:
|
||||
output: quant-dequant weight
|
||||
"""
|
||||
org_shape = data.shape
|
||||
weight, scale, zp = quant_tensor(data, num_bits, group_size, scheme, dtype, ratio)
|
||||
return np.reshape(scale * (weight - zp), org_shape)
|
||||
|
||||
|
||||
def pad_tensor(weight, group_size, k_blocks):
|
||||
"""Pad tensor rowi so that it can be is divisible by group_size.
|
||||
|
||||
Args:
|
||||
weight (array): weight
|
||||
group_size (int): how many elements share one scale/zp
|
||||
k_blocks (int): the number of block
|
||||
|
||||
Returns:
|
||||
weight: paded weight
|
||||
"""
|
||||
if group_size == -1:
|
||||
return weight
|
||||
|
||||
org_w_shape = weight.shape
|
||||
padded_rows = k_blocks * group_size
|
||||
pad_len = padded_rows - org_w_shape[0]
|
||||
|
||||
if pad_len > 0:
|
||||
weight = np.pad(weight, ((0, pad_len), (0, 0)), "constant")
|
||||
|
||||
return weight
|
||||
|
||||
|
||||
def rtn_quantize(
|
||||
model,
|
||||
weight_config={}, # noqa: B006
|
||||
num_bits=4,
|
||||
group_size=32,
|
||||
scheme="asym",
|
||||
ratios={}, # noqa: B006
|
||||
accuracy_level=0,
|
||||
providers=["CPUExecutionProvider"], # noqa: B006
|
||||
algorithm="k_quant",
|
||||
):
|
||||
"""Quant the model with round to nearst method.
|
||||
|
||||
Args:
|
||||
model (ModelProto or ONNXModel): onnx model
|
||||
weight_config (dict): quantization config
|
||||
For example,
|
||||
weight_config = {
|
||||
'fc2':
|
||||
{
|
||||
'bits': 4,
|
||||
'group_size': 32,
|
||||
'scheme': 'sym',
|
||||
'algorithm': 'RTN'
|
||||
}
|
||||
}
|
||||
num_bits (int, optional): num_bits. Default is 4.
|
||||
group_size (int, optional): how many elements share one scale/zp. Default is 32.
|
||||
scheme (str, optional): sym or asym. Defaults to "asym".
|
||||
ratios (dict, optional): percentile of clip. Defaults to {}.
|
||||
accuracy_level (int): accuracy level. Support 0 (unset),1(fp32), 2(fp16), 3(bf16), or 4(int8).
|
||||
providers (list): providers to use
|
||||
|
||||
Returns:
|
||||
model: fake quantized ONNXModel
|
||||
"""
|
||||
model = ONNXModel(model)
|
||||
base_dir = os.path.dirname(model.model_path) if model.model_path is not None else ""
|
||||
new_nodes = []
|
||||
remove_nodes = []
|
||||
total_num = len([i for i in model.nodes() if i.op_type in ["MatMul"]])
|
||||
curr_id = 0
|
||||
for node in model.nodes():
|
||||
if node.op_type in ["MatMul"]:
|
||||
curr_id += 1
|
||||
simple_progress_bar(total_num, curr_id)
|
||||
if (
|
||||
node.op_type in ["MatMul"]
|
||||
and model.get_initializer(node.input[1]) is not None
|
||||
and weight_config.get(node.name, {}) != "fp32"
|
||||
):
|
||||
weight_tensor = model.get_initializer(node.input[1])
|
||||
weight = numpy_helper.to_array(weight_tensor, base_dir=base_dir).copy()
|
||||
if len(weight.shape) != 2:
|
||||
continue
|
||||
|
||||
dtype = weight.dtype
|
||||
|
||||
if node.name in weight_config:
|
||||
num_bits = weight_config[node.name]["bits"]
|
||||
group_size = weight_config[node.name]["group_size"]
|
||||
scheme = weight_config[node.name]["scheme"]
|
||||
|
||||
org_w_shape = weight.shape # ic, oc
|
||||
group_size = group_size if group_size != -1 else org_w_shape[0]
|
||||
|
||||
k_blocks = (org_w_shape[0] - 1) // group_size + 1
|
||||
init_share_num = model.get_initializer_share_num(node.input[1])
|
||||
|
||||
weight = pad_tensor(weight, group_size, k_blocks)
|
||||
|
||||
satisfy_MatMulNBits_condition = num_bits == 4 or num_bits == 8 # noqa: N806
|
||||
|
||||
if satisfy_MatMulNBits_condition: # pragma: no cover
|
||||
if algorithm == "k_quant":
|
||||
q_weight, scale, zp = quant_tensor_k_quant_cuda(weight.T, num_bits, group_size)
|
||||
else:
|
||||
q_weight, scale, zp = quant_tensor(
|
||||
weight.T, num_bits, group_size, scheme, "uint", ratios.get(node.input[1], 1)
|
||||
)
|
||||
|
||||
q_matmul_node, new_inits = make_matmul_weight_only_node(
|
||||
node=node,
|
||||
weight_shape=org_w_shape,
|
||||
num_bits=num_bits,
|
||||
group_size=group_size,
|
||||
k_blocks=k_blocks,
|
||||
q_weight=q_weight.astype("uint8"),
|
||||
scale=scale.astype(dtype),
|
||||
zero_point=zp if scheme == "asym" or algorithm == "k_quant" else None,
|
||||
accuracy_level=accuracy_level,
|
||||
)
|
||||
|
||||
model.add_initializers(new_inits)
|
||||
remove_nodes.append(node)
|
||||
new_nodes.append(q_matmul_node)
|
||||
else:
|
||||
q_weight = qdq_tensor(weight.T, num_bits, group_size, scheme, "int", ratios.get(node.input[1], 1))
|
||||
q_weight = np.reshape(q_weight, (org_w_shape[1], -1))
|
||||
q_weight = np.transpose(q_weight)
|
||||
q_weight = q_weight[: org_w_shape[0], :].astype(dtype)
|
||||
q_weight_tensor = onnx.helper.make_tensor(
|
||||
name=node.input[1] + f"_Q{num_bits!s}G{group_size!s}",
|
||||
data_type=np_dtype_to_tensor_dtype(dtype),
|
||||
dims=weight.shape,
|
||||
vals=q_weight.tobytes(),
|
||||
raw=True,
|
||||
)
|
||||
model.add_initializer(q_weight_tensor)
|
||||
node.input[1] = q_weight_tensor.name
|
||||
if init_share_num == 1:
|
||||
model.remove_initializer(weight_tensor)
|
||||
|
||||
model.add_nodes(new_nodes)
|
||||
model.remove_nodes(remove_nodes)
|
||||
model.topological_sort()
|
||||
return model
|
||||
|
||||
|
||||
def get_weight_scale(weight, group_size):
|
||||
"""Get the scale of weight."""
|
||||
org_shape = weight.shape
|
||||
weight = np.reshape(weight, (-1, group_size)) if group_size != -1 else weight
|
||||
scale = np.mean(np.reshape(np.abs(weight) / np.max(np.abs(weight), axis=1, keepdims=True), org_shape), axis=0)
|
||||
return scale
|
||||
|
||||
|
||||
def prepare_inputs(model, n_samples, dataloader, providers):
|
||||
"""Prepare inputs for weight only quantization.
|
||||
|
||||
Args:
|
||||
model (ModelProto or ONNXModel): onnx model
|
||||
n_samples (int, optional): calibration sample number. -1 means all samples.
|
||||
dataloader (object): dataloader for calibration.
|
||||
providers (list): providers to use
|
||||
|
||||
Returns:
|
||||
inputs: prepared inputs.
|
||||
so: session options
|
||||
"""
|
||||
from importlib.util import find_spec # noqa: PLC0415
|
||||
|
||||
from .util import to_numpy # noqa: PLC0415
|
||||
|
||||
so = ort.SessionOptions()
|
||||
if sys.version_info < (3, 11) and find_spec("onnxruntime_extensions"): # pragma: no cover
|
||||
from onnxruntime_extensions import get_library_path # noqa: PLC0415
|
||||
|
||||
so.register_custom_ops_library(get_library_path())
|
||||
if model.is_large_model:
|
||||
onnx.save_model(
|
||||
model.model,
|
||||
model.model_path + "_augment.onnx",
|
||||
save_as_external_data=True,
|
||||
all_tensors_to_one_file=True,
|
||||
convert_attribute=False,
|
||||
)
|
||||
|
||||
session = (
|
||||
ort.InferenceSession(model.model.SerializeToString(), so, providers=providers)
|
||||
if not model.is_large_model
|
||||
else ort.InferenceSession(model.model_path + "_augment.onnx", so, providers=providers)
|
||||
)
|
||||
inputs_names = [i.name for i in session.get_inputs()]
|
||||
del session
|
||||
|
||||
inputs = []
|
||||
for i, data in enumerate(dataloader):
|
||||
if n_samples != -1 and ((i + 1) * dataloader.batch_size) > n_samples:
|
||||
break
|
||||
if len(inputs_names) != 1 or isinstance(data[0], dict):
|
||||
assert len(data[0]) == len(inputs_names), (
|
||||
f"Input number mismatch, require {len(inputs_names)} but get {len(data[0])}"
|
||||
)
|
||||
|
||||
if isinstance(data[0], dict):
|
||||
inputs.append(dict([(name, to_numpy(inp_data)) for name, inp_data in data[0].items()])) # noqa: C404
|
||||
elif isinstance(data[0], np.ndarray): # pragma: no cover
|
||||
inputs.append(dict([(name, inp) for name, inp in zip(inputs_names, [data[0]], strict=False)])) # noqa: C404
|
||||
else: # pragma: no cover
|
||||
inputs.append(dict([(name, to_numpy(inp)) for name, inp in zip(inputs_names, data[0], strict=False)])) # noqa: C404
|
||||
return inputs, so
|
||||
|
||||
|
||||
def gptq(
|
||||
W,
|
||||
H,
|
||||
num_bits=4,
|
||||
group_size=32,
|
||||
scheme="asym",
|
||||
blocksize=128,
|
||||
percdamp=0.01,
|
||||
actorder=False,
|
||||
mse=False,
|
||||
perchannel=True,
|
||||
):
|
||||
"""Quant the weight with GPTQ method.
|
||||
|
||||
Args:
|
||||
W (array): weight.
|
||||
H (array): Hessian matrix.
|
||||
num_bits (int, optional): num_bits. Default is 4.
|
||||
group_size (int, optional): how many elements share one scale/zp. Default is 32.
|
||||
scheme (str, optional): sym or asym. Defaults to "asym".
|
||||
blocksize (int, optional): blocksize to quantize weight.
|
||||
percdamp (float, optional): percent of the average Hessian diagonal to use for dampening.
|
||||
actorder (bool, optional): whether rearrange Hessian matrix considering the diag's value.
|
||||
mse (bool, optional): whether get scale and zero point with mse error.
|
||||
perchannel (bool, optional): whether quantize weight per-channel.
|
||||
|
||||
Returns:
|
||||
Q: fake quantized weight
|
||||
"""
|
||||
maxq = 2**num_bits - 1
|
||||
grid = 100
|
||||
maxshrink = 0.8
|
||||
norm = 2.4
|
||||
|
||||
def find_params(weight):
|
||||
org_shape = weight.shape
|
||||
# find zp, scale
|
||||
if not perchannel:
|
||||
weight = np.expand_dims(weight.flatten(), axis=1)
|
||||
tmp = np.zeros(weight.shape[1])
|
||||
xmin = np.minimum(np.min(weight, axis=0), tmp)
|
||||
xmax = np.maximum(np.max(weight, axis=0), tmp)
|
||||
if scheme == "sym":
|
||||
xmax = np.maximum(np.abs(xmin), xmax)
|
||||
tmp = xmin < 0
|
||||
if np.any(tmp):
|
||||
xmin[tmp] = -xmax[tmp]
|
||||
tmp = (xmin == 0) & (xmax == 0)
|
||||
xmin[tmp] = -1
|
||||
xmax[tmp] = +1
|
||||
|
||||
scale = (xmax - xmin) / maxq
|
||||
if scheme == "sym":
|
||||
zero = np.ones(scale.shape) * (maxq + 1) / 2
|
||||
else:
|
||||
zero = np.round(-xmin / scale)
|
||||
if mse:
|
||||
best = np.ones([weight.shape[1]]) * float("inf")
|
||||
for i in range(int(maxshrink * grid)):
|
||||
p = 1 - i / grid
|
||||
xmin1 = p * xmin
|
||||
xmax1 = p * xmax
|
||||
scale1 = (xmax1 - xmin1) / maxq
|
||||
zero1 = np.round(-xmin1 / scale1) if scheme != "sym" else zero
|
||||
q = np.clip(np.round(weight / scale1) + zero1, 0, maxq)
|
||||
q -= weight
|
||||
q = np.power(np.abs(q), norm)
|
||||
err = np.sum(q, 0)
|
||||
tmp = err < best
|
||||
if np.any(tmp):
|
||||
best[tmp] = err[tmp]
|
||||
scale[tmp] = scale1[tmp]
|
||||
zero[tmp] = zero1[tmp]
|
||||
if not perchannel:
|
||||
tmp = org_shape[1]
|
||||
scale = np.repeat(scale, tmp)
|
||||
zero = np.repeat(zero, tmp)
|
||||
shape = [-1] + [1] * (len(org_shape) - 1)
|
||||
scale = np.reshape(scale, shape)
|
||||
zero = np.reshape(zero, shape)
|
||||
return scale, zero
|
||||
|
||||
shape = W.shape
|
||||
scale, zp = find_params(W)
|
||||
dead = np.diag(H) == 0
|
||||
H[dead, dead] = 1
|
||||
W[dead, :] = 0 # such channel makes no contribution to quantization computation
|
||||
|
||||
# rearrange considering the diag's value
|
||||
if actorder:
|
||||
perm = np.argsort(np.diag(H))[::-1]
|
||||
W = W[perm, :] # noqa: N806
|
||||
H = H[perm, :][:, perm] # noqa: N806
|
||||
Losses = np.zeros_like(W) # noqa: N806
|
||||
Q = np.zeros_like(W) # noqa: N806
|
||||
damp = percdamp * np.mean(np.diag(H))
|
||||
diag = np.arange(shape[0])
|
||||
H[diag, diag] += damp # add a average value of
|
||||
H = np.linalg.cholesky(np.linalg.inv(H)).T # noqa: N806
|
||||
Hinv = H # noqa: N806
|
||||
for i1 in range(0, shape[0], blocksize):
|
||||
i2 = min(i1 + blocksize, shape[0])
|
||||
count = i2 - i1
|
||||
|
||||
W1 = copy.deepcopy(W[i1:i2, :]) # noqa: N806
|
||||
Q1 = np.zeros_like(W1) # noqa: N806
|
||||
Err1 = np.zeros_like(W1) # noqa: N806
|
||||
Losses1 = np.zeros_like(W1) # noqa: N806
|
||||
Hinv1 = Hinv[i1:i2, i1:i2] # noqa: N806
|
||||
|
||||
for i in range(count): # within a block, channel wise
|
||||
w = W1[i, :]
|
||||
d = Hinv1[i, i]
|
||||
|
||||
if group_size != -1:
|
||||
if (i1 + i) % group_size == 0:
|
||||
scale, zp = find_params(W[(i1 + i) : (i1 + i + group_size), :])
|
||||
|
||||
q = (scale * (np.clip(np.round(w[:, np.newaxis] / scale) + zp, 0, maxq) - zp)).flatten()
|
||||
Q1[i, :] = q
|
||||
Losses1[i, :] = (w - q) ** 2 / d**2
|
||||
|
||||
err1 = (w - q) / d
|
||||
W1[i:, :] -= np.matmul(np.expand_dims(Hinv1[i:, i], axis=1), np.expand_dims(err1, axis=0))
|
||||
Err1[i, :] = err1
|
||||
|
||||
Q[i1:i2, :] = Q1
|
||||
Losses[i1:i2, :] = Losses1 / 2
|
||||
|
||||
W[i2:, :] -= np.matmul(Hinv[i2:, i1:i2], Err1)
|
||||
|
||||
if actorder:
|
||||
invperm = np.argsort(perm)
|
||||
Q = Q[invperm, :] # noqa: N806
|
||||
|
||||
Q = np.reshape(Q, W.shape) # noqa: N806
|
||||
del W
|
||||
return Q
|
||||
|
||||
|
||||
def gptq_quantize(
|
||||
model,
|
||||
dataloader,
|
||||
weight_config={}, # noqa: B006
|
||||
num_bits=4,
|
||||
group_size=32,
|
||||
scheme="asym",
|
||||
n_samples=128,
|
||||
percdamp=0.01,
|
||||
blocksize=128,
|
||||
actorder=False,
|
||||
mse=False,
|
||||
perchannel=True,
|
||||
accuracy_level=0,
|
||||
providers=["CPUExecutionProvider"], # noqa: B006
|
||||
):
|
||||
"""Quant the model with GPTQ method.
|
||||
|
||||
Args:
|
||||
model (ModelProto or ONNXModel): onnx model
|
||||
dataloader (object): dataloader for calibration.
|
||||
weight_config (dict): quantization config
|
||||
For example,
|
||||
weight_config = {
|
||||
'fc2':
|
||||
{
|
||||
'bits': 4,
|
||||
'group_size': 32,
|
||||
'scheme': 'sym',
|
||||
'algorithm': 'GPTQ'
|
||||
}
|
||||
}
|
||||
num_bits (int, optional): num_bits. Default is 4.
|
||||
group_size (int, optional): how many elements share one scale/zp. Default is 32.
|
||||
scheme (str, optional): sym or asym. Defaults to "asym".
|
||||
n_samples (int, optional): calibration sample number.
|
||||
percdamp (float, optional): percent of the average Hessian diagonal to use for dampening.
|
||||
blocksize (int, optional): blocksize to quantize weight.
|
||||
actorder (bool, optional): whether rearrange Hessian matrix considering the diag's value.
|
||||
mse (bool, optional): whether get scale and zero point with mse error.
|
||||
perchannel (bool, optional): whether quantize weight per-channel.
|
||||
accuracy_level (int): accuracy level. Support 0 (unset), 1(fp32), 2(fp16), 3(bf16), or 4(int8).
|
||||
providers (list): providers to use
|
||||
|
||||
Returns:
|
||||
model: fake quantized ONNXModel
|
||||
"""
|
||||
model = ONNXModel(model)
|
||||
base_dir = os.path.dirname(model.model_path) if model.model_path is not None else ""
|
||||
|
||||
inputs, so = prepare_inputs(model, n_samples, dataloader, providers)
|
||||
del dataloader
|
||||
org_output = copy.deepcopy(model.model.graph.output)
|
||||
model.remove_tensors_from_outputs([i.name for i in org_output])
|
||||
output_names = []
|
||||
for node in model.nodes():
|
||||
if (
|
||||
node.op_type in ["MatMul"]
|
||||
and weight_config.get(node.name, {}) != "fp32"
|
||||
and weight_config.get(node.name, {}).get("algorithm", "GPTQ") == "GPTQ"
|
||||
):
|
||||
output_names.append(node.input[0])
|
||||
output_names = list(set(output_names))
|
||||
model.add_tensors_to_outputs(output_names)
|
||||
if model.is_large_model:
|
||||
onnx.save_model(
|
||||
model.model,
|
||||
model.model_path + "_augment.onnx",
|
||||
save_as_external_data=True,
|
||||
all_tensors_to_one_file=True,
|
||||
convert_attribute=False,
|
||||
)
|
||||
|
||||
session = (
|
||||
ort.InferenceSession(model.model.SerializeToString(), so, providers=providers)
|
||||
if not model.is_large_model
|
||||
else ort.InferenceSession(model.model_path + "_augment.onnx", so, providers=providers)
|
||||
)
|
||||
|
||||
for idx, input_name in enumerate(output_names):
|
||||
simple_progress_bar(len(output_names), idx + 1)
|
||||
node_list = []
|
||||
weights = []
|
||||
|
||||
for node in model.input_name_to_nodes[input_name]:
|
||||
if (
|
||||
node.op_type in ["MatMul"]
|
||||
and weight_config.get(node.name, {}) != "fp32"
|
||||
and weight_config.get(node.name, {}).get("algorithm", "GPTQ") == "GPTQ"
|
||||
and model.get_initializer(node.input[1]) is not None
|
||||
):
|
||||
weight = numpy_helper.to_array(
|
||||
model.get_initializer(model.get_node(node.name).input[1]), base_dir
|
||||
).copy()
|
||||
if len(weight.shape) != 2:
|
||||
continue
|
||||
|
||||
weights.append(weight)
|
||||
node_list.append(model.get_node(node.name))
|
||||
|
||||
if len(weights) == 0:
|
||||
continue
|
||||
|
||||
Hs = [np.zeros((i.shape[0], i.shape[0])) for i in weights] # noqa: N806
|
||||
nsamples = 0
|
||||
for data in inputs:
|
||||
inp = session.run([input_name], data)[0]
|
||||
tmp = inp.shape[0]
|
||||
inp = np.reshape(inp, (-1, inp.shape[-1]))
|
||||
Hs = [i * (nsamples / (nsamples + tmp)) for i in Hs] # noqa: N806
|
||||
nsamples += tmp
|
||||
inp = np.sqrt(2 / nsamples) * inp
|
||||
Hs = [i + np.matmul(inp.T, inp) for i in Hs] # noqa: N806
|
||||
|
||||
for (
|
||||
node,
|
||||
weight,
|
||||
H, # noqa: N806
|
||||
) in zip(node_list, weights, Hs, strict=False):
|
||||
if node.name in weight_config:
|
||||
num_bits = weight_config[node.name]["bits"]
|
||||
group_size = weight_config[node.name]["group_size"]
|
||||
scheme = weight_config[node.name]["scheme"]
|
||||
group_size = group_size if group_size != -1 else weight.shape[0]
|
||||
dtype = weight.dtype
|
||||
|
||||
q_weight = gptq(
|
||||
weight,
|
||||
H,
|
||||
num_bits=num_bits,
|
||||
group_size=group_size,
|
||||
scheme=scheme,
|
||||
blocksize=blocksize,
|
||||
percdamp=percdamp,
|
||||
actorder=actorder,
|
||||
mse=mse,
|
||||
perchannel=perchannel,
|
||||
)
|
||||
|
||||
weight_tensor = model.get_initializer(node.input[1])
|
||||
init_share_num = model.get_initializer_share_num(node.input[1])
|
||||
|
||||
satisfy_MatMulNBits_condition = num_bits == 4 # noqa: N806
|
||||
|
||||
if satisfy_MatMulNBits_condition: # pragma: no cover
|
||||
org_shape = weight.shape
|
||||
k_blocks = (org_shape[0] + group_size - 1) // group_size
|
||||
q_weight = pad_tensor(q_weight, group_size, k_blocks)
|
||||
q_weight, scale, zp = quant_tensor(q_weight.T, num_bits, group_size, scheme, "uint")
|
||||
q_matmul_node, new_inits = make_matmul_weight_only_node(
|
||||
node=node,
|
||||
weight_shape=org_shape,
|
||||
num_bits=num_bits,
|
||||
group_size=group_size,
|
||||
k_blocks=k_blocks,
|
||||
q_weight=q_weight.astype("uint8"),
|
||||
scale=scale.astype(dtype),
|
||||
zero_point=zp if scheme == "asym" else None,
|
||||
accuracy_level=accuracy_level,
|
||||
)
|
||||
|
||||
model.add_initializers(new_inits)
|
||||
model.remove_node(node)
|
||||
model.add_node(q_matmul_node)
|
||||
else:
|
||||
q_weight_tensor = onnx.helper.make_tensor(
|
||||
name=node.input[1] + f"_Q{num_bits!s}G{group_size!s}",
|
||||
data_type=np_dtype_to_tensor_dtype(dtype),
|
||||
dims=q_weight.shape,
|
||||
vals=q_weight.astype(dtype).tobytes(),
|
||||
raw=True,
|
||||
)
|
||||
model.add_initializer(q_weight_tensor)
|
||||
node.input[1] = q_weight_tensor.name
|
||||
if init_share_num == 1:
|
||||
model.remove_initializer(weight_tensor)
|
||||
|
||||
model.remove_tensors_from_outputs(output_names)
|
||||
model.model.graph.output.MergeFrom(org_output)
|
||||
|
||||
model.topological_sort()
|
||||
|
||||
# reload external data to prevent external data file path errors
|
||||
if model.is_large_model:
|
||||
from onnx.external_data_helper import load_external_data_for_model # noqa: PLC0415
|
||||
|
||||
load_external_data_for_model(model.model, os.path.split(model.model_path)[0])
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,600 @@
|
||||
# --------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
from pathlib import Path
|
||||
|
||||
import onnx
|
||||
import onnx.helper as onnx_helper
|
||||
import onnx.numpy_helper as onnx_numpy_helper
|
||||
from onnx.onnx_pb import ModelProto
|
||||
|
||||
from .quant_utils import attribute_to_kwarg, find_by_name
|
||||
|
||||
|
||||
def _clean_initializers_helper(graph, model):
|
||||
"""Clean unused initializers from graph.
|
||||
|
||||
Returns:
|
||||
A cleaned graph without unused initializers
|
||||
A list of tensor names, which are not produced by this graph and its subgraphes
|
||||
"""
|
||||
requesting_tensor_names = set()
|
||||
requesting_tensor_names.update(input_name for node in graph.node for input_name in node.input if input_name)
|
||||
requesting_tensor_names.update(g_out.name for g_out in graph.output if g_out.name)
|
||||
|
||||
new_nodes = []
|
||||
for node in graph.node:
|
||||
new_node = node
|
||||
graph_attrs = [
|
||||
attr
|
||||
for attr in node.attribute
|
||||
if attr.type == onnx.AttributeProto.GRAPH or attr.type == onnx.AttributeProto.GRAPHS
|
||||
]
|
||||
if graph_attrs:
|
||||
kwargs = {}
|
||||
for attr in node.attribute:
|
||||
new_attribute = {}
|
||||
if attr.type == onnx.AttributeProto.GRAPH:
|
||||
(
|
||||
cleaned_sub_graph,
|
||||
sub_requesting_tensor_names,
|
||||
) = _clean_initializers_helper(attr.g, model)
|
||||
new_attribute = {attr.name: cleaned_sub_graph}
|
||||
requesting_tensor_names.update(sub_requesting_tensor_names)
|
||||
elif attr.type == onnx.AttributeProto.GRAPHS:
|
||||
cleaned_graphes = []
|
||||
for subgraph in attr.graphs:
|
||||
(
|
||||
cleaned_sub_graph,
|
||||
sub_requesting_tensor_names,
|
||||
) = _clean_initializers_helper(subgraph, model)
|
||||
cleaned_graphes.append(cleaned_sub_graph)
|
||||
requesting_tensor_names.update(sub_requesting_tensor_names)
|
||||
new_attribute = {attr.name: cleaned_graphes}
|
||||
else:
|
||||
new_attribute = attribute_to_kwarg(attr)
|
||||
kwargs.update(new_attribute)
|
||||
new_node = onnx_helper.make_node(node.op_type, node.input, node.output, name=node.name, **kwargs)
|
||||
new_nodes.append(new_node)
|
||||
|
||||
graph.ClearField("node")
|
||||
graph.node.extend(new_nodes)
|
||||
|
||||
requesting_tensor_names.difference_update(output for node in graph.node for output in node.output)
|
||||
|
||||
unused_initializer = []
|
||||
for initializer in graph.initializer:
|
||||
if initializer.name in requesting_tensor_names:
|
||||
requesting_tensor_names.remove(initializer.name)
|
||||
else:
|
||||
# mark it to remove, remove here directly will cause mis-behavier
|
||||
unused_initializer.append(initializer)
|
||||
|
||||
name_to_input = {input.name: input for input in graph.input}
|
||||
for initializer in unused_initializer:
|
||||
graph.initializer.remove(initializer)
|
||||
if initializer.name in name_to_input:
|
||||
try:
|
||||
graph.input.remove(name_to_input[initializer.name])
|
||||
except StopIteration:
|
||||
if model.ir_version < 4:
|
||||
print(f"Warning: invalid weight name {initializer.name} found in the graph (not a graph input)")
|
||||
|
||||
requesting_tensor_names.difference_update(input.name for input in graph.input)
|
||||
|
||||
return graph, requesting_tensor_names
|
||||
|
||||
|
||||
class ONNXModel:
|
||||
def __init__(self, model: ModelProto):
|
||||
self.model = model
|
||||
|
||||
def nodes(self):
|
||||
return self.model.graph.node
|
||||
|
||||
def initializer(self):
|
||||
return self.model.graph.initializer
|
||||
|
||||
def initializer_extend(self, inits):
|
||||
if len(inits) == 0:
|
||||
raise ValueError("Can add an empty list.")
|
||||
for init in self.initializer():
|
||||
self._check_init(init, "gain")
|
||||
for init in inits:
|
||||
self._check_init(init)
|
||||
self.model.graph.initializer.append(init)
|
||||
|
||||
def graph(self):
|
||||
return self.model.graph
|
||||
|
||||
def ir_version(self):
|
||||
return self.model.ir_version
|
||||
|
||||
def opset_import(self):
|
||||
return self.model.opset_import
|
||||
|
||||
def set_opset_import(self, domain, version):
|
||||
for opset in self.model.opset_import:
|
||||
if opset.domain == domain:
|
||||
opset.version = version
|
||||
return
|
||||
|
||||
self.model.opset_import.extend([onnx_helper.make_opsetid(domain, version)])
|
||||
|
||||
def remove_node(self, node):
|
||||
if node in self.model.graph.node:
|
||||
self.model.graph.node.remove(node)
|
||||
|
||||
def remove_nodes(self, nodes_to_remove):
|
||||
for node in nodes_to_remove:
|
||||
self.remove_node(node)
|
||||
|
||||
def add_node(self, node):
|
||||
self.model.graph.node.extend([self._check_node(node)])
|
||||
|
||||
def add_nodes(self, nodes_to_add):
|
||||
for node in nodes_to_add:
|
||||
self.add_node(node)
|
||||
|
||||
def add_initializer(self, tensor):
|
||||
if find_by_name(tensor.name, self.model.graph.initializer) is None:
|
||||
self._check_init(tensor)
|
||||
self.model.graph.initializer.extend([tensor])
|
||||
|
||||
def get_initializer(self, name):
|
||||
for tensor in self.model.graph.initializer:
|
||||
if tensor.name == name:
|
||||
return tensor
|
||||
return None
|
||||
|
||||
def find_graph_input(self, input_name):
|
||||
for input in self.model.graph.input:
|
||||
if input.name == input_name:
|
||||
return input
|
||||
return None
|
||||
|
||||
def find_graph_output(self, output_name):
|
||||
for output in self.model.graph.output:
|
||||
if output.name == output_name:
|
||||
return output
|
||||
return None
|
||||
|
||||
def get_tensor_type(self, tensor_name: str):
|
||||
tensor_type_map = {obj.name: obj.type for obj in self.model.graph.value_info}
|
||||
|
||||
if tensor_name in tensor_type_map:
|
||||
return tensor_type_map[tensor_name].tensor_type
|
||||
|
||||
g_input = self.find_graph_input(tensor_name)
|
||||
if g_input:
|
||||
return g_input.type.tensor_type
|
||||
|
||||
g_output = self.find_graph_output(tensor_name)
|
||||
if g_output:
|
||||
return g_output.type.tensor_type
|
||||
|
||||
return None
|
||||
|
||||
def get_constant_value(self, output_name):
|
||||
for node in self.model.graph.node:
|
||||
if node.op_type == "Constant":
|
||||
if node.output[0] == output_name:
|
||||
for attr in node.attribute:
|
||||
if attr.name == "value":
|
||||
return onnx_numpy_helper.to_array(attr.t)
|
||||
|
||||
# Fallback to initializer since constant folding may have been applied.
|
||||
initializer = self.get_initializer(output_name)
|
||||
if initializer is not None:
|
||||
return onnx_numpy_helper.to_array(initializer)
|
||||
|
||||
return None
|
||||
|
||||
def get_initializer_name_set(self):
|
||||
return {initializer.name for initializer in self.model.graph.initializer}
|
||||
|
||||
def remove_initializer(self, tensor):
|
||||
if tensor in self.model.graph.initializer:
|
||||
self.model.graph.initializer.remove(tensor)
|
||||
for input in self.model.graph.input:
|
||||
if input.name == tensor.name:
|
||||
self.model.graph.input.remove(input)
|
||||
break
|
||||
|
||||
def remove_initializers(self, init_to_remove):
|
||||
for initializer in init_to_remove:
|
||||
self.remove_initializer(initializer)
|
||||
|
||||
def get_non_initializer_inputs(self):
|
||||
initializer_names = self.get_initializer_name_set()
|
||||
non_initializer_inputs = set()
|
||||
for input in self.model.graph.input:
|
||||
if input.name not in initializer_names:
|
||||
non_initializer_inputs.add(input.name)
|
||||
return non_initializer_inputs
|
||||
|
||||
def input_name_to_nodes(self):
|
||||
input_name_to_nodes = {}
|
||||
for node in self.model.graph.node:
|
||||
for input_name in node.input:
|
||||
if input_name: # Could be empty when it is optional
|
||||
if input_name not in input_name_to_nodes:
|
||||
input_name_to_nodes[input_name] = [node]
|
||||
else:
|
||||
input_name_to_nodes[input_name].append(node)
|
||||
return input_name_to_nodes
|
||||
|
||||
def output_name_to_node(self):
|
||||
output_name_to_node = {}
|
||||
for node in self.model.graph.node:
|
||||
for output_name in node.output:
|
||||
if output_name: # Could be empty when it is optional
|
||||
output_name_to_node[output_name] = node
|
||||
return output_name_to_node
|
||||
|
||||
def get_children(self, node, input_name_to_nodes=None):
|
||||
if input_name_to_nodes is None:
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
|
||||
children = []
|
||||
for output in node.output:
|
||||
if output in input_name_to_nodes:
|
||||
for node in input_name_to_nodes[output]:
|
||||
children.append(node) # noqa: PERF402
|
||||
return children
|
||||
|
||||
def get_parents(self, node, output_name_to_node=None):
|
||||
if output_name_to_node is None:
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
||||
parents = []
|
||||
for input in node.input:
|
||||
if input in output_name_to_node:
|
||||
parents.append(output_name_to_node[input])
|
||||
return parents
|
||||
|
||||
def get_parent(self, node, idx, output_name_to_node=None):
|
||||
if output_name_to_node is None:
|
||||
output_name_to_node = self.output_name_to_node()
|
||||
|
||||
if len(node.input) <= idx:
|
||||
return None
|
||||
|
||||
input = node.input[idx]
|
||||
if input not in output_name_to_node:
|
||||
return None
|
||||
|
||||
return output_name_to_node[input]
|
||||
|
||||
def find_node_by_name(self, node_name, new_nodes_list, graph):
|
||||
"""Find out if a node exists in a graph or a node is in the
|
||||
new set of nodes created during quantization.
|
||||
|
||||
Returns:
|
||||
The node found or None.
|
||||
"""
|
||||
graph_nodes_list = list(graph.node) # deep copy
|
||||
graph_nodes_list.extend(new_nodes_list)
|
||||
node = find_by_name(node_name, graph_nodes_list)
|
||||
return node
|
||||
|
||||
def get_largest_node_name_suffix(self, node_name_prefix):
|
||||
"""
|
||||
Gets the largest node name (int) suffix for all node names that begin with `node_name_prefix`.
|
||||
Example: for nodes my_prefix_0 and my_prefix_3, this method returns 3.
|
||||
"""
|
||||
suffix = -1
|
||||
|
||||
for node in self.model.graph.node:
|
||||
if node.name and node.name.startswith(node_name_prefix):
|
||||
try:
|
||||
index = int(node.name[len(node_name_prefix) :])
|
||||
suffix = max(index, suffix)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return suffix
|
||||
|
||||
def get_largest_initializer_name_suffix(self, initializer_name_prefix):
|
||||
"""
|
||||
Gets the largest initializer name integer suffix for all initializer names that begin
|
||||
with `initializer_name_prefix`. This can be used to create unique initializer names.
|
||||
|
||||
Example: for initializer names 'my_weight_0' and 'my_weight_3', this method returns 3 if
|
||||
`initializer_name_prefix` is 'my_weight_'.
|
||||
"""
|
||||
suffix = -1
|
||||
|
||||
for initializer in self.model.graph.initializer:
|
||||
if initializer.name.startswith(initializer_name_prefix):
|
||||
try:
|
||||
index = int(initializer.name[len(initializer_name_prefix) :])
|
||||
suffix = max(index, suffix)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return suffix
|
||||
|
||||
def find_nodes_by_initializer(self, graph, initializer):
|
||||
"""
|
||||
Find all nodes with given initializer as an input.
|
||||
"""
|
||||
nodes = []
|
||||
for node in graph.node:
|
||||
for node_input in node.input:
|
||||
if node_input == initializer.name:
|
||||
nodes.append(node)
|
||||
return nodes
|
||||
|
||||
@staticmethod
|
||||
def __get_initializer(name, graph_path):
|
||||
for gid in range(len(graph_path) - 1, -1, -1):
|
||||
graph = graph_path[gid]
|
||||
for tensor in graph.initializer:
|
||||
if tensor.name == name:
|
||||
return tensor, graph
|
||||
return None, None
|
||||
|
||||
@staticmethod
|
||||
def __replace_gemm_with_matmul(graph_path):
|
||||
new_nodes = []
|
||||
graph = graph_path[-1]
|
||||
for node in graph.node:
|
||||
graph_attrs = [attr for attr in node.attribute if attr.type == 5 or attr.type == 10]
|
||||
if graph_attrs:
|
||||
kwargs = {}
|
||||
for attr in node.attribute:
|
||||
if attr.type == 5:
|
||||
graph_path.append(attr.g)
|
||||
kv = {attr.name: ONNXModel.__replace_gemm_with_matmul(graph_path)}
|
||||
elif attr.type == 10:
|
||||
value = []
|
||||
for subgraph in attr.graphs:
|
||||
graph_path.append(subgraph)
|
||||
value.extend([ONNXModel.__replace_gemm_with_matmul(graph_path)])
|
||||
kv = {attr.name: value}
|
||||
else:
|
||||
kv = attribute_to_kwarg(attr)
|
||||
kwargs.update(kv)
|
||||
node = onnx_helper.make_node( # noqa: PLW2901
|
||||
node.op_type, node.input, node.output, name=node.name, **kwargs
|
||||
)
|
||||
|
||||
if node.op_type == "Gemm":
|
||||
alpha = 1.0
|
||||
beta = 1.0
|
||||
transA = 0 # noqa: N806
|
||||
transB = 0 # noqa: N806
|
||||
for attr in node.attribute:
|
||||
if attr.name == "alpha":
|
||||
alpha = onnx_helper.get_attribute_value(attr)
|
||||
elif attr.name == "beta":
|
||||
beta = onnx_helper.get_attribute_value(attr)
|
||||
elif attr.name == "transA":
|
||||
transA = onnx_helper.get_attribute_value(attr) # noqa: N806
|
||||
elif attr.name == "transB":
|
||||
transB = onnx_helper.get_attribute_value(attr) # noqa: N806
|
||||
if alpha == 1.0 and beta == 1.0 and transA == 0:
|
||||
inputB = node.input[1] # noqa: N806
|
||||
if transB == 1:
|
||||
B, Bs_graph = ONNXModel.__get_initializer(node.input[1], graph_path) # noqa: N806
|
||||
if B:
|
||||
# assume B is not used by any other node
|
||||
B_array = onnx_numpy_helper.to_array(B) # noqa: N806
|
||||
B_trans = onnx_numpy_helper.from_array(B_array.T) # noqa: N806
|
||||
B_trans.name = B.name
|
||||
Bs_graph.initializer.remove(B)
|
||||
for input in Bs_graph.input:
|
||||
if input.name == inputB:
|
||||
Bs_graph.input.remove(input)
|
||||
break
|
||||
Bs_graph.initializer.extend([B_trans])
|
||||
else:
|
||||
inputB += "_Transposed" # noqa: N806
|
||||
transpose_node = onnx_helper.make_node(
|
||||
"Transpose",
|
||||
inputs=[node.input[1]],
|
||||
outputs=[inputB],
|
||||
name=node.name + "_Transpose" if node.name else "",
|
||||
)
|
||||
new_nodes.append(transpose_node)
|
||||
|
||||
matmul_node = onnx_helper.make_node(
|
||||
"MatMul",
|
||||
inputs=[node.input[0], inputB],
|
||||
outputs=[node.output[0] + ("_MatMul" if len(node.input) > 2 else "")],
|
||||
name=node.name + "_MatMul" if node.name else "",
|
||||
)
|
||||
new_nodes.append(matmul_node)
|
||||
|
||||
if len(node.input) > 2:
|
||||
add_node = onnx_helper.make_node(
|
||||
"Add",
|
||||
inputs=[node.output[0] + "_MatMul", node.input[2]],
|
||||
outputs=node.output,
|
||||
name=node.name + "_Add" if node.name else "",
|
||||
)
|
||||
new_nodes.append(add_node)
|
||||
|
||||
# unsupported
|
||||
else:
|
||||
new_nodes.append(node)
|
||||
|
||||
# not GEMM
|
||||
else:
|
||||
new_nodes.append(node)
|
||||
|
||||
graph.ClearField("node")
|
||||
graph.node.extend(new_nodes)
|
||||
graph_path.pop()
|
||||
return graph
|
||||
|
||||
def replace_gemm_with_matmul(self):
|
||||
graph_path = [self.graph()]
|
||||
ONNXModel.__replace_gemm_with_matmul(graph_path)
|
||||
|
||||
def save_model_to_file(self, output_path, use_external_data_format=False):
|
||||
"""
|
||||
Save model to external data, which is needed for model size > 2GB
|
||||
"""
|
||||
self.topological_sort()
|
||||
if use_external_data_format:
|
||||
onnx.external_data_helper.convert_model_to_external_data(
|
||||
self.model,
|
||||
all_tensors_to_one_file=True,
|
||||
location=Path(output_path).name + ".data",
|
||||
convert_attribute=True,
|
||||
)
|
||||
for init in self.model.graph.initializer:
|
||||
self._check_init(init, "end")
|
||||
onnx.save_model(self.model, output_path)
|
||||
|
||||
@staticmethod
|
||||
def replace_node_input(node, old_input_name, new_input_name):
|
||||
assert isinstance(old_input_name, str) and isinstance(new_input_name, str)
|
||||
for j in range(len(node.input)):
|
||||
if node.input[j] == old_input_name:
|
||||
node.input[j] = new_input_name
|
||||
|
||||
def replace_input_of_all_nodes(self, old_input_name, new_input_name):
|
||||
for node in self.model.graph.node:
|
||||
ONNXModel.replace_node_input(node, old_input_name, new_input_name)
|
||||
|
||||
def replace_input_of_nodes(self, old_input_name, new_input_name, node_names_set):
|
||||
for node in self.model.graph.node:
|
||||
if node.name in node_names_set:
|
||||
ONNXModel.replace_node_input(node, old_input_name, new_input_name)
|
||||
|
||||
@staticmethod
|
||||
def replace_node_output(node, old_output_name, new_output_name):
|
||||
assert isinstance(old_output_name, str) and isinstance(new_output_name, str)
|
||||
for j in range(len(node.output)):
|
||||
if node.output[j] == old_output_name:
|
||||
node.output[j] = new_output_name
|
||||
|
||||
def replace_output_of_all_nodes(self, old_output_name, new_output_name):
|
||||
for node in self.model.graph.node:
|
||||
ONNXModel.replace_node_output(node, old_output_name, new_output_name)
|
||||
|
||||
def replace_output_of_nodes(self, old_output_name, new_output_name, node_names_set):
|
||||
for node in self.model.graph.node:
|
||||
if node.name in node_names_set:
|
||||
ONNXModel.replace_node_output(node, old_output_name, new_output_name)
|
||||
|
||||
def remove_unused_constant(self):
|
||||
input_name_to_nodes = self.input_name_to_nodes()
|
||||
|
||||
# remove unused constant
|
||||
unused_nodes = []
|
||||
nodes = self.nodes()
|
||||
for node in nodes:
|
||||
if (
|
||||
node.op_type == "Constant"
|
||||
and not self.is_graph_output(node.output[0])
|
||||
and node.output[0] not in input_name_to_nodes
|
||||
):
|
||||
unused_nodes.append(node)
|
||||
|
||||
self.remove_nodes(unused_nodes)
|
||||
|
||||
ununsed_weights = []
|
||||
for w in self.initializer():
|
||||
if w.name not in input_name_to_nodes and not self.is_graph_output(w.name):
|
||||
ununsed_weights.append(w)
|
||||
# Remove from graph.input
|
||||
for graph_input in self.graph().input:
|
||||
if graph_input.name == w.name:
|
||||
self.graph().input.remove(graph_input)
|
||||
|
||||
self.remove_initializers(ununsed_weights)
|
||||
|
||||
def is_graph_output(self, output_name):
|
||||
return any(output.name == output_name for output in self.model.graph.output)
|
||||
|
||||
def is_graph_input(self, tensor_name: str) -> bool:
|
||||
return any(input.name == tensor_name for input in self.model.graph.input)
|
||||
|
||||
# TODO:use OnnxModel.graph_topological_sort(self.model.graph) from transformers.onnx_model
|
||||
# Currently it breaks Openvino/Linux training gpu pipeline so hold off for 1.8 release
|
||||
def topological_sort(self):
|
||||
deps_count = [0] * len(self.nodes()) # dependency count of each node
|
||||
deps_to_nodes = {} # input to node indice
|
||||
sorted_nodes = [] # initialize sorted_nodes
|
||||
for node_idx, node in enumerate(self.nodes()):
|
||||
# CANNOT use len(node.input) directly because input can be optional
|
||||
deps_count[node_idx] = sum(1 for _ in node.input if _)
|
||||
if deps_count[node_idx] == 0: # Constant doesn't depend on any inputs
|
||||
sorted_nodes.append(self.nodes()[node_idx])
|
||||
continue
|
||||
|
||||
for input_name in node.input:
|
||||
if not input_name:
|
||||
continue
|
||||
if input_name not in deps_to_nodes:
|
||||
deps_to_nodes[input_name] = [node_idx]
|
||||
else:
|
||||
deps_to_nodes[input_name].append(node_idx)
|
||||
|
||||
initializer_names = [init.name for init in self.initializer()]
|
||||
graph_input_names = [input.name for input in self.model.graph.input]
|
||||
input_names = initializer_names + graph_input_names
|
||||
input_names.sort()
|
||||
prev_input_name = None
|
||||
for input_name in input_names:
|
||||
if prev_input_name == input_name:
|
||||
continue
|
||||
|
||||
prev_input_name = input_name
|
||||
if input_name in deps_to_nodes:
|
||||
for node_idx in deps_to_nodes[input_name]:
|
||||
deps_count[node_idx] = deps_count[node_idx] - 1
|
||||
if deps_count[node_idx] == 0:
|
||||
sorted_nodes.append(self.nodes()[node_idx])
|
||||
|
||||
start = 0
|
||||
end = len(sorted_nodes)
|
||||
|
||||
while start < end:
|
||||
for output in sorted_nodes[start].output:
|
||||
if output in deps_to_nodes:
|
||||
for node_idx in deps_to_nodes[output]:
|
||||
deps_count[node_idx] = deps_count[node_idx] - 1
|
||||
if deps_count[node_idx] == 0:
|
||||
sorted_nodes.append(self.nodes()[node_idx])
|
||||
end = end + 1
|
||||
start = start + 1
|
||||
|
||||
assert end == len(self.graph().node), "Graph is not a DAG"
|
||||
self.graph().ClearField("node")
|
||||
self.graph().node.extend(sorted_nodes)
|
||||
|
||||
def clean_initializers(self):
|
||||
return _clean_initializers_helper(self.graph(), self.model)
|
||||
|
||||
def _check_init(self, init, test=None):
|
||||
if init.data_type == onnx.TensorProto.FLOAT8E4M3FN:
|
||||
if init.HasField("raw_data"):
|
||||
b = list(init.raw_data)
|
||||
if any((i & 127) == 127 for i in b):
|
||||
raise ValueError(f"Initializer {init.name!r} has nan.")
|
||||
return init
|
||||
|
||||
def _check_node(self, node):
|
||||
"""
|
||||
A quantization to float 8 does not use quantized bias but float 16 bias.
|
||||
This function checks that DequantizeLinear is not used to
|
||||
dequantize from float 16.
|
||||
"""
|
||||
if node.op_type == "DequantizeLinear":
|
||||
zero_point = node.input[2]
|
||||
init = self.get_initializer(zero_point)
|
||||
dtype = init.data_type
|
||||
if dtype in {
|
||||
onnx.TensorProto.FLOAT16,
|
||||
onnx.TensorProto.FLOAT,
|
||||
onnx.TensorProto.DOUBLE,
|
||||
onnx.TensorProto.BFLOAT16,
|
||||
}:
|
||||
raise RuntimeError(f"Unsupported DequantizeLinear operator, dequantization from {dtype}.")
|
||||
return node
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,2 @@
|
||||
# from .base_operator import QuantOperatorBase
|
||||
# from .matmul import MatMulInteger
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user