转换

量化转换/恢复工具。

函数

replace_quant_module

递归地将模块替换为量化模块。

set_quantizer_by_cfg

根据指定的quant_cfg更新量化器属性。

set_quantizer_attribute

通过通配符或过滤函数对量化器属性进行细粒度调整。

register

为给定的未量化原始类注册一个量化类。

unregister

取消为给定的未量化原始类注册的量化类。

set_quantizer_by_cfg_context

用于使用quant_cfg设置量化器属性的上下文管理器。
register(original_cls, quantized_cls)

为给定的未量化原始类注册一个量化类。

Parameters:

  • original_cls (Module) – 原始的未量化类。

  • quantized_cls (Module) – 量化类。该类应具有一个_setup方法，该方法初始化在前向传播中调用的各种量化器。量化类的前向传播函数应在正确的位置调用量化器。

以下是一个定义量化类并注册它的示例：

import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.nn import TensorQuantizer


class QuantLayerNorm(nn.LayerNorm):
    def __init__(self, normalized_shape):
        super().__init__(normalized_shape)
        self._setup()

    def _setup(self):
        # Method to setup the quantizers
        self.input_quantizer = TensorQuantizer()
        self.weight_quantizer = TensorQuantizer()

    def forward(self, input):
        input = self.input_quantizer(input)
        weight = self.weight_quantizer(self.weight)
        return F.layer_norm(input, self.normalized_shape, weight, self.bias, self.eps)


# Register the custom quantized module
mtq.register(original_cls=nn.LayerNorm, quantized_cls=QuantLayerNorm)
replace_quant_module(model, version=None)

递归地将模块替换为量化模块。

Parameters:

模型 (模块) –

set_quantizer_attribute(quant_model, wildcard_or_filter_func, attribute, parent_class=None)

通过通配符或过滤函数对量化器属性进行细粒度调整。

Parameters:

  • quant_model (Module) – 一个pytorch模型

  • wildcard_or_filter_func (str | Callable) – 一个通配符字符串或过滤函数。通配符字符串与量化器模块名称匹配。量化器模块是 TensorQuantizer的实例。 过滤函数以量化模块名称为输入，如果应调整量化器则返回True，否则返回False。

  • attribute (QuantizerAttributeConfig | List[QuantizerAttributeConfig] | Dict[str | Callable, QuantizerAttributeConfig | List[QuantizerAttributeConfig]] | Dict | List[Dict]) – An instance of QuantizerAttributeConfig or an equivalent dictionary or a list of these two types. If attribute is a list, the matched TensorQuantizer modules will be replaced with SequentialQuantizer modules having one quantizer for each attribute instance from the list. See set_from_attribute_config() for more details on the supported attributes and their types.

  • parent_class (None | type) – （可选）与wildcard_or_filter_func匹配的量化器模块的父类，应该进行调整。如果None，所有匹配的量化器模块都将被调整。

set_quantizer_by_cfg(quant_model, quant_cfg)

根据指定的quant_cfg更新量化器属性。

quant_cfg 是一个字典，将通配符或过滤函数映射到其量化器属性，这些属性在 QuantizerAttributeConfig 中定义。通配符或过滤函数与量化器模块名称进行匹配。匹配的量化器模块的指定量化器属性将相应设置。键 "default" 是一个特殊键，用于设置所有量化器的量化器属性，这些量化器的模块名称没有匹配到其他通配符或过滤函数。

此外，字典条目也可以是映射到特定类量化配置的pytorch模块类名。pytorch模块应该有一个量化等效项。

有关更多详细信息，请参见set_quantizer_attribute。

Parameters:
set_quantizer_by_cfg_context(quant_model, quant_cfg)

用于使用quant_cfg设置量化器属性的上下文管理器。

退出上下文管理器后，设置的属性将被重置为原始属性。 有关更多详细信息，请参见set_quantizer_by_cfg()。

请谨慎使用此上下文管理器。更改量化器的某些属性，例如calibrator，可能会导致意外行为。

Parameters:
unregister(original_cls)

取消为给定的未量化原始类注册的量化类。

Parameters:

original_cls (Module) – 原始的未量化类。