mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
56 lines
2.4 KiB
Python
56 lines
2.4 KiB
Python
from fnmatch import fnmatch
|
|
from typing import Any, Dict, List, Optional, Union
|
|
import torch
|
|
|
|
from optimum.quanto.quantize import _quantize_submodule
|
|
from optimum.quanto.tensor import Optimizer, qtype
|
|
|
|
# the quantize function in quanto had a bug where it was using exclude instead of include
|
|
|
|
|
|
def quantize(
|
|
model: torch.nn.Module,
|
|
weights: Optional[Union[str, qtype]] = None,
|
|
activations: Optional[Union[str, qtype]] = None,
|
|
optimizer: Optional[Optimizer] = None,
|
|
include: Optional[Union[str, List[str]]] = None,
|
|
exclude: Optional[Union[str, List[str]]] = None,
|
|
):
|
|
"""Quantize the specified model submodules
|
|
|
|
Recursively quantize the submodules of the specified parent model.
|
|
|
|
Only modules that have quantized counterparts will be quantized.
|
|
|
|
If include patterns are specified, the submodule name must match one of them.
|
|
|
|
If exclude patterns are specified, the submodule must not match one of them.
|
|
|
|
Include or exclude patterns are Unix shell-style wildcards which are NOT regular expressions. See
|
|
https://docs.python.org/3/library/fnmatch.html for more details.
|
|
|
|
Note: quantization happens in-place and modifies the original model and its descendants.
|
|
|
|
Args:
|
|
model (`torch.nn.Module`): the model whose submodules will be quantized.
|
|
weights (`Optional[Union[str, qtype]]`): the qtype for weights quantization.
|
|
activations (`Optional[Union[str, qtype]]`): the qtype for activations quantization.
|
|
include (`Optional[Union[str, List[str]]]`):
|
|
Patterns constituting the allowlist. If provided, module names must match at
|
|
least one pattern from the allowlist.
|
|
exclude (`Optional[Union[str, List[str]]]`):
|
|
Patterns constituting the denylist. If provided, module names must not match
|
|
any patterns from the denylist.
|
|
"""
|
|
if include is not None:
|
|
include = [include] if isinstance(include, str) else include
|
|
if exclude is not None:
|
|
exclude = [exclude] if isinstance(exclude, str) else exclude
|
|
for name, m in model.named_modules():
|
|
if include is not None and not any(fnmatch(name, pattern) for pattern in include):
|
|
continue
|
|
if exclude is not None and any(fnmatch(name, pattern) for pattern in exclude):
|
|
continue
|
|
_quantize_submodule(model, name, m, weights=weights,
|
|
activations=activations, optimizer=optimizer)
|