pytorch-quantization 的文档
基本功能
量化函数
tensor_quant 和 fake_tensor_quant 是两个量化张量的基本函数。fake_tensor_quant 返回伪量化张量（浮点值）。tensor_quant 返回量化张量（整数值）和比例。
tensor_quant(inputs, amax, num_bits=8, output_dtype=torch.float, unsigned=False)
fake_tensor_quant(inputs, amax, num_bits=8, output_dtype=torch.float, unsigned=False)
示例
from pytorch_quantization import tensor_quant
# Generate random input. With fixed seed 12345, x should be
# tensor([0.9817, 0.8796, 0.9921, 0.4611, 0.0832, 0.1784, 0.3674, 0.5676, 0.3376, 0.2119])
torch.manual_seed(12345)
x = torch.rand(10)
# fake quantize tensor x. fake_quant_x will be
# tensor([0.9843, 0.8828, 0.9921, 0.4609, 0.0859, 0.1797, 0.3672, 0.5703, 0.3359, 0.2109])
fake_quant_x = tensor_quant.fake_tensor_quant(x, x.abs().max())
# quantize tensor x. quant_x will be
# tensor([126., 113., 127., 59., 11., 23., 47., 73., 43., 27.])
# with scale=128.0057
quant_x, scale = tensor_quant.tensor_quant(x, x.abs().max())
两个函数的反向传播都定义为 Straight-Through Estimator (STE)。
描述符和量化器
QuantDescriptor 定义了张量应该如何量化。还有一些预定义的 QuantDescriptor，例如 QUANT_DESC_8BIT_PER_TENSOR 和 QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL。
TensorQuantizer 是用于量化张量的模块，由 QuantDescriptor 定义。
from pytorch_quantization.tensor_quant import QuantDescriptor
from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
quant_desc = QuantDescriptor(num_bits=4, fake_quant=False, axis=(0), unsigned=True)
quantizer = TensorQuantizer(quant_desc)
torch.manual_seed(12345)
x = torch.rand(10, 9, 8, 7)
quant_x = quantizer(x)
如果在 QuantDescriptor 中给出了 amax，TensorQuantizer 将使用它进行量化。否则，TensorQuantizer 将计算 amax 然后进行量化。amax 将根据指定的 axis 计算。请注意，QuantDescriptor 的 axis 指定剩余轴，与 max() 的轴相反。
量化模块
主要有两种类型的模块，Conv 和 Linear。两者都可以替代 torch.nn 版本，并在权重和激活上应用量化。
除了原始模块的参数外，两者都接受 quant_desc_input 和 quant_desc_weight。
from torch import nn
from pytorch_quantization import tensor_quant
import pytorch_quantization.nn as quant_nn
# pytorch's module
fc1 = nn.Linear(in_features, out_features, bias=True)
conv1 = nn.Conv2d(in_channels, out_channels, kernel_size)
# quantized version
quant_fc1 = quant_nn.Linear(
in_features, out_features, bias=True,
quant_desc_input=tensor_quant.QUANT_DESC_8BIT_PER_TENSOR,
quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW)
quant_conv1 = quant_nn.Conv2d(
in_channels, out_channels, kernel_size,
quant_desc_input=tensor_quant.QUANT_DESC_8BIT_PER_TENSOR,
quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)
训练后量化
一个模型可以通过简单地调用 quant_modules.initialize() 来进行训练后量化
from pytorch_quantization import quant_modules
model = torchvision.models.resnet50()
如果模型不是完全由模块定义的，则应该手动创建 TensorQuantizer 并将其添加到模型中的正确位置。
校准
校准是 TensorRT 的术语，指的是将数据样本传递给量化器，并确定激活的最佳 amax。我们支持 3 种校准方法
max：简单地使用全局最大绝对值entropy：TensorRT 的熵校准percentile：根据给定的百分位数去除异常值。mse：基于 MSE（均方误差）的校准
在上面的 ResNet50 示例中，校准方法设置为 mse，可以按以下示例使用
# Find the TensorQuantizer and enable calibration
for name, module in model.named_modules():
if name.endswith('_quantizer'):
module.enable_calib()
module.disable_quant() # Use full precision data to calibrate
# Feeding data samples
model(x)
# ...
# Finalize calibration
for name, module in model.named_modules():
if name.endswith('_quantizer'):
module.load_calib_amax()
module.disable_calib()
module.enable_quant()
# If running on GPU, it needs to call .cuda() again because new tensors will be created by calibration process
model.cuda()
# Keep running the quantized model
# ...
注意
校准需要在将模型导出到 ONNX 之前执行。
量化感知训练
量化感知训练基于 Straight Through Estimator (STE) 导数近似。它有时被称为“量化感知训练”。我们不使用这个名称，因为它没有反映底层的假设。如果说有什么不同的话，那就是由于 STE 近似，它使训练“不感知”量化。
完成校准后，量化感知训练只需选择一个训练计划并继续训练校准后的模型。通常，不需要微调很长时间。我们通常使用原始训练计划的 10% 左右，从初始训练学习率的 1% 开始，并使用余弦退火学习率计划，该计划遵循余弦周期下降的一半，降至初始微调学习率的 1%（初始训练学习率的 0.01%）。
一些建议
量化感知训练（本质上是一个离散数值优化问题）在数学上不是一个已解决的问题。根据我们的经验，以下是一些建议
为了使 STE 近似效果良好，最好使用较小的学习率。较大的学习率更可能扩大 STE 近似引入的方差，并破坏训练好的网络。
在训练期间不要更改量化表示（比例），至少不要太频繁。每步都更改比例，实际上就像每步都更改数据格式（e8m7、e5m10、e3m4 等），这将很容易影响收敛。
导出到 ONNX
导出到 ONNX 的目标是部署到 TensorRT，而不是 ONNX 运行时。因此，我们只将伪量化模型导出为 TensorRT 可以接受的形式。伪量化将被分解为一对 QuantizeLinear/DequantizeLinear ONNX 算子。TensorRT 将接受生成的 ONNX 图，并以其能力范围内最优化方式在 int8 中执行它。
注意
目前，我们只支持导出 int8 和 fp8 伪量化模块。此外，量化模块需要在导出到 ONNX 之前进行校准。
伪量化模型可以像任何其他 Pytorch 模型一样导出到 ONNX。请在 torch.onnx 了解更多关于将 Pytorch 模型导出到 ONNX 的信息。例如
import pytorch_quantization
from pytorch_quantization import nn as quant_nn
from pytorch_quantization import quant_modules
quant_modules.initialize()
model = torchvision.models.resnet50()
# load the calibrated model
state_dict = torch.load("quant_resnet50-entropy-1024.pth", map_location="cpu")
model.load_state_dict(state_dict)
model.cuda()
dummy_input = torch.randn(128, 3, 224, 224, device='cuda')
input_names = [ "actual_input_1" ]
output_names = [ "output1" ]
with pytorch_quantization.enable_onnx_export():
# enable_onnx_checker needs to be disabled. See notes below.
torch.onnx.export(
model, dummy_input, "quant_resnet50.onnx", verbose=True, opset_version=10, enable_onnx_checker=False
)
注意
请注意，在 opset13 中，axis 已添加到 QuantizeLinear 和 DequantizeLinear 中。
量化 Resnet50
创建量化模型
导入必要的 python 模块
import torch
import torch.utils.data
from torch import nn
from pytorch_quantization import nn as quant_nn
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor
from torchvision import models
sys.path.append("path to torchvision/references/classification/")
from train import evaluate, train_one_epoch, load_data
添加量化模块
第一步是将量化器模块添加到神经网络图中。此包提供了许多量化层模块，其中包含用于输入和权重的量化器。例如 quant_nn.QuantLinear，可以代替 nn.Linear 使用。这些量化层可以通过猴子补丁自动替换，也可以通过手动修改模型定义来替换。
自动层替换通过 quant_modules 完成。这应该在模型创建之前调用。
from pytorch_quantization import quant_modules
quant_modules.initialize()
这将应用于每个模块的所有实例。如果您不希望所有模块都被量化，则应该手动替换量化模块。独立的量化器也可以使用 quant_nn.TensorQuantizer 添加到模型中。
训练后量化
为了实现高效的推理，我们希望为每个量化器选择一个固定的范围。从预训练模型开始，最简单的方法是通过校准。
校准
我们将对激活使用基于直方图的校准，对权重使用默认的最大值校准。
quant_desc_input = QuantDescriptor(calib_method='histogram')
quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
model = models.resnet50(pretrained=True)
model.cuda()
为了收集激活直方图，我们必须将样本数据输入到模型中。首先，创建 ImageNet 数据加载器，就像在训练脚本中所做的那样。然后，在每个量化器中启用校准，并将训练数据输入到模型中。1024 个样本（2 批 512 个）应该足以估计激活的分布。使用训练数据进行校准，以便验证也测量所选范围的泛化能力。
data_path = "PATH to imagenet"
batch_size = 512
traindir = os.path.join(data_path, 'train')
valdir = os.path.join(data_path, 'val')
dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, False, False)
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size,
sampler=train_sampler, num_workers=4, pin_memory=True)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=batch_size,
sampler=test_sampler, num_workers=4, pin_memory=True)
def collect_stats(model, data_loader, num_batches):
"""Feed data to the network and collect statistic"""
# Enable calibrators
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
module.disable_quant()
module.enable_calib()
else:
module.disable()
for i, (image, _) in tqdm(enumerate(data_loader), total=num_batches):
model(image.cuda())
if i >= num_batches:
break
# Disable calibrators
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
module.enable_quant()
module.disable_calib()
else:
module.enable()
def compute_amax(model, **kwargs):
# Load calib result
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
if isinstance(module._calibrator, calib.MaxCalibrator):
module.load_calib_amax()
else:
module.load_calib_amax(**kwargs)
print(F"{name:40}: {module}")
model.cuda()
# It is a bit slow since we collect histograms on CPU
with torch.no_grad():
collect_stats(model, data_loader, num_batches=2)
compute_amax(model, method="percentile", percentile=99.99)
完成校准后，量化器将设置 amax，它表示量化空间中可表示的最大绝对输入值。默认情况下，权重范围是按通道的，而激活范围是按张量的。我们可以通过打印每个 TensorQuantizer 模块来查看压缩的 amax。
conv1._input_quantizer : TensorQuantizer(8bit fake per-tensor amax=2.6400 calibrator=MaxCalibrator(track_amax=False) quant)
conv1._weight_quantizer : TensorQuantizer(8bit fake axis=(0) amax=[0.0000, 0.7817](64) calibrator=MaxCalibrator(track_amax=False) quant)
layer1.0.conv1._input_quantizer : TensorQuantizer(8bit fake per-tensor amax=6.8645 calibrator=MaxCalibrator(track_amax=False) quant)
layer1.0.conv1._weight_quantizer : TensorQuantizer(8bit fake axis=(0) amax=[0.0000, 0.7266](64) calibrator=MaxCalibrator(track_amax=False) quant)
...
评估校准后的模型
接下来，我们将评估训练后量化模型在 ImageNet 验证集上的分类准确率。
criterion = nn.CrossEntropyLoss()
with torch.no_grad():
evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)
# Save the model
torch.save(model.state_dict(), "/tmp/quant_resnet50-calibrated.pth")
这应该产生 76.1% 的 top-1 准确率，这与预训练模型的 76.2% 准确率接近。
使用不同的校准
我们可以尝试不同的校准方法，而无需重新收集直方图，看看哪种方法能获得最佳准确率。
with torch.no_grad():
compute_amax(model, method="percentile", percentile=99.9)
evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)
with torch.no_grad():
for method in ["mse", "entropy"]:
print(F"{method} calibration")
compute_amax(model, method=method)
evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)
MSE 和熵都应该超过 76%。99.9% 的百分位数对于 resnet50 来说裁剪了太多的值，并且会获得稍低的准确率。
量化感知训练
可选地，我们可以微调校准后的模型以进一步提高准确率。
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
# Training takes about one and half hour per epoch on a single V100
train_one_epoch(model, criterion, optimizer, data_loader, "cuda", 0, 100)
# Save the model
torch.save(model.state_dict(), "/tmp/quant_resnet50-finetuned.pth")
经过一个 epoch 的微调后，我们可以达到超过 76.4% 的 top-1 准确率。使用学习率退火进行更多 epoch 的微调可以进一步提高准确率。例如，使用余弦退火以 0.001 的学习率开始微调 15 个 epoch 可以获得超过 76.7% 的准确率。应该注意的是，相同的微调计划也将提高未量化模型的准确率。
进一步优化
为了在 TensorRT 上实现高效的推理，我们需要了解更多关于运行时优化的细节。TensorRT 支持量化卷积和残差加法的融合。新的融合算子有两个输入。让我们称它们为 conv-input 和 residual-input。这里融合算子的输出精度必须与残差输入精度匹配。当融合算子之后有另一个量化节点时，我们可以在残差输入和 Elementwise-Addition 节点之间插入一对量化/反量化节点，以便卷积节点之后的量化节点与卷积节点融合，并且卷积节点完全使用 INT8 输入和输出进行量化。我们不能使用自动猴子补丁来应用此优化，我们需要手动插入量化/反量化节点。
首先从 https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py 复制 resnet.py，修改构造函数，添加显式布尔标志 ‘quantize’
def resnet50(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, quantize, **kwargs)
def _resnet(arch: str, block: Type[Union[BasicBlock, Bottleneck]], layers: List[int], pretrained: bool, progress: bool,
quantize: bool, **kwargs: Any) -> ResNet:
model = ResNet(block, layers, quantize, **kwargs)
class ResNet(nn.Module):
def __init__(self,
block: Type[Union[BasicBlock, Bottleneck]],
layers: List[int],
quantize: bool = False,
num_classes: int = 1000,
zero_init_residual: bool = False,
groups: int = 1,
width_per_group: int = 64,
replace_stride_with_dilation: Optional[List[bool]] = None,
norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
super(ResNet, self).__init__()
self._quantize = quantize
当此 self._quantize 标志设置为 True 时，我们需要将所有 nn.Conv2d 替换为 quant_nn.QuantConv2d。
def conv3x3(in_planes: int,
out_planes: int,
stride: int = 1,
groups: int = 1,
dilation: int = 1,
quantize: bool = False) -> nn.Conv2d:
"""3x3 convolution with padding"""
if quantize:
return quant_nn.QuantConv2d(in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=dilation,
groups=groups,
bias=False,
dilation=dilation)
else:
return nn.Conv2d(in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=dilation,
groups=groups,
bias=False,
dilation=dilation)
def conv1x1(in_planes: int, out_planes: int, stride: int = 1, quantize: bool = False) -> nn.Conv2d:
"""1x1 convolution"""
if quantize:
return quant_nn.QuantConv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
else:
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
残差 conv 加法可以在 BasicBlock 和 Bottleneck 中找到。我们需要首先在 __init__ 函数中声明量化节点。
def __init__(self,
inplanes: int,
planes: int,
stride: int = 1,
downsample: Optional[nn.Module] = None,
groups: int = 1,
base_width: int = 64,
dilation: int = 1,
norm_layer: Optional[Callable[..., nn.Module]] = None,
quantize: bool = False) -> None:
# other code...
self._quantize = quantize
if self._quantize:
self.residual_quantizer = quant_nn.TensorQuantizer(quant_nn.QuantConv2d.default_quant_desc_input)
最后，我们需要修补 BasicBlock 和 Bottleneck 中的 forward 函数，在此处插入额外的量化/反量化节点。
def forward(self, x: Tensor) -> Tensor:
# other code...
if self._quantize:
out += self.residual_quantizer(identity)
else:
out += identity
out = self.relu(out)
return out
带有残差量化的最终 resnet 代码可以在 https://github.com/NVIDIA/TensorRT/blob/master/tools/pytorch-quantization/examples/torchvision/models/classification/resnet.py 中找到
创建自定义量化模块
量化工具提供了几个量化模块，如下所示
QuantConv1d, QuantConv2d, QuantConv3d, QuantConvTranspose1d, QuantConvTranspose2d, QuantConvTranspose3d
QuantLinear
QuantAvgPool1d, QuantAvgPool2d, QuantAvgPool3d, QuantMaxPool1d, QuantMaxPool2d, QuantMaxPool3d
要量化一个模块，我们需要量化输入和权重（如果存在）。以下是 3 个主要用例
为只有输入的模块创建量化包装器
为既有输入又有权重的模块创建量化包装器。
直接将
TensorQuantizer模块添加到模型图中操作的输入。
如果需要自动将原始模块（图中的节点）替换为其量化版本，则前两种方法非常有用。当需要在非常特定的位置手动将量化添加到模型图时，第三种方法可能很有用（更手动，更多控制）。
让我们通过下面的示例来了解每个用例。
量化只有输入的模块
一个合适的例子是量化 pooling 模块变体。
本质上，我们需要提供一个包装器函数，该函数接受原始模块并在其周围添加 TensorQuantizer 模块，以便首先量化输入，然后将其馈送到原始模块中。
通过子类化原始模块 (
pooling.MaxPool2d) 以及实用程序模块 (_utils.QuantInputMixin) 来创建包装器。
class QuantMaxPool2d(pooling.MaxPool2d, _utils.QuantInputMixin):
__init__.py函数将调用原始模块的 init 函数，并为其提供相应的参数。将只有一个使用**kwargs的附加参数，其中包含量化配置信息。QuantInputMixin实用程序包含方法pop_quant_desc_in_kwargs，该方法从输入中提取此配置信息，如果输入为None，则返回默认值。最后，调用init_quantizer方法，该方法初始化TensorQuantizer模块，该模块将量化输入。
def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
return_indices=False, ceil_mode=False, **kwargs):
super(QuantMaxPool2d, self).__init__(kernel_size, stride, padding, dilation,
return_indices, ceil_mode)
quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
self.init_quantizer(quant_desc_input)
初始化后，需要在我们的包装器模块中定义
forward函数，该函数将实际使用在__init__函数中初始化的_input_quantizer量化输入，并使用super调用将输入转发到基础模块。
def forward(self, input):
quant_input = self._input_quantizer(input)
return super(QuantMaxPool2d, self).forward(quant_input)
最后，我们需要为
_input_quantizer定义一个 getter 方法。例如，这可以用于禁用特定模块的量化，使用module.input_quantizer.disable()，这在试验不同的层量化配置时很有用。
@property
def input_quantizer(self):
return self._input_quantizer
一个完整的量化池化模块如下所示
class QuantMaxPool2d(pooling.MaxPool2d, _utils.QuantInputMixin):
"""Quantized 2D maxpool"""
def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
return_indices=False, ceil_mode=False, **kwargs):
super(QuantMaxPool2d, self).__init__(kernel_size, stride, padding, dilation,
return_indices, ceil_mode)
quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
self.init_quantizer(quant_desc_input)
def forward(self, input):
quant_input = self._input_quantizer(input)
return super(QuantMaxPool2d, self).forward(quant_input)
@property
def input_quantizer(self):
return self._input_quantizer
量化带有权重和输入的模块
我们给出了量化 torch.nn.Linear 模块的示例。由此可见，与之前的量化池化模块的示例相比，唯一额外的更改是我们需要适应 Linear 模块中权重的量化。
我们创建量化的线性模块如下
class QuantLinear(nn.Linear, _utils.QuantMixin):
在
__init__函数中，我们首先使用pop_quant_desc_in_kwargs函数提取输入和权重的量化描述符。其次，我们使用这些量化描述符初始化输入和权重的TensorQuantizer模块。
def __init__(self, in_features, out_features, bias=True, **kwargs):
super(QuantLinear, self).__init__(in_features, out_features, bias)
quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
self.init_quantizer(quant_desc_input, quant_desc_weight)
此外，覆盖
forward函数调用，并在将量化参数传递给实际的F.Linear调用之前，分别通过_input_quantizer和_weight_quantizer传递输入和权重。此步骤将实际的输入/权重TensorQuantizer添加到模块，最终添加到模型。
def forward(self, input):
quant_input = self._input_quantizer(input)
quant_weight = self._weight_quantizer(self.weight)
output = F.linear(quant_input, quant_weight, bias=self.bias)
return output
与
Linear模块类似，我们为与输入/权重关联的TensorQuantizer模块添加了 getter 方法。例如，这可以用于通过调用module_obj.weight_quantizer.disable()来禁用量化机制
@property
def input_quantizer(self):
return self._input_quantizer
@property
def weight_quantizer(self):
return self._weight_quantizer
通过以上所有更改，量化的 Linear 模块将如下所示
class QuantLinear(nn.Linear, _utils.QuantMixin):
def __init__(self, in_features, out_features, bias=True, **kwargs):
super(QuantLinear, self).__init__(in_features, out_features, bias)
quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
self.init_quantizer(quant_desc_input, quant_desc_weight)
def forward(self, input):
quant_input = self._input_quantizer(input)
quant_weight = self._weight_quantizer(self.weight)
output = F.linear(quant_input, quant_weight, bias=self.bias)
return output
@property
def input_quantizer(self):
return self._input_quantizer
@property
def weight_quantizer(self):
return self._weight_quantizer
直接量化图中的输入
也可以直接量化图输入，而无需像上面解释的那样创建包装器。
这是一个例子
test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
quantizer = TensorQuantizer(quant_nn.QuantLinear.default_quant_desc_input)
quant_input = quantizer(test_input)
out = F.adaptive_avg_pool2d(quant_input, 3)
假设图中有一个 F.adaptive_avg_pool2d 操作，我们想量化此操作。在上面的示例中，我们使用 TensorQuantizer(quant_nn.QuantLinear.default_quant_desc_input) 定义一个量化器，然后我们使用该量化器来实际量化 test_input，然后将此量化输入馈送到 F.adaptive_avg_pool2d 操作。请注意，此量化器与我们之前在创建 torch 模块的量化版本时使用的量化器相同。
pytorch_quantization.calib
pytorch_quantization.calib 提供了 Calibrator 类，用于收集数据统计信息并确定 pytorch_quantization 参数。
MaxCalibrator
- class pytorch_quantization.calib.MaxCalibrator(num_bits, axis, unsigned, track_amax=False)
最大值校准器，全局跟踪最大值
- 参数
calib_desc – MaxCalibDescriptor。
num_bits – 整数。量化位数。
axis – 元组。参见 QuantDescriptor。
unsigned – 布尔值。使用无符号量化。
- 只读属性
amaxs：amax 列表。Numpy 数组被保存，因为它可能用于某些绘图。
- collect(x)
跟踪所有张量的绝对最大值
- 参数
x – 张量
- Raises
RuntimeError – 如果 amax 形状更改
- compute_amax()
返回收集的所有张量的绝对最大值
- reset()
重置收集的绝对最大值
HistogramCalibrator
- class pytorch_quantization.calib.HistogramCalibrator(num_bits, axis, unsigned, num_bins=2048, grow_method=None, skip_zeros=False, torch_hist=True)
统一的直方图校准器
- 直方图将只收集一次。compute_amax() 基于参数执行熵、百分位数或 mse 校准
基于参数的校准
- 参数
num_bits – 整数。量化位数。
axis – 元组。参见 QuantDescriptor。
unsigned – 布尔值。使用无符号量化。
num_bins – 整数。直方图箱的数量。默认为 2048。
grow_method – 字符串。已弃用。默认为 None。
skip_zeros – 布尔值。如果为 True，则在收集直方图数据时跳过零。默认为 False。
torch_hist – 布尔值。如果为 True，则通过 torch.histc 而不是 np.histogram 收集直方图。如果输入张量在 GPU 上，则 histc 也将在 GPU 上运行。默认为 True。
- collect(x)
收集直方图
- compute_amax(method: str, *, stride: int = 1, start_bin: int = 128, percentile: float = 99.99)
从收集的直方图计算 amax
- 参数
method – 字符串。以下之一 [‘entropy’, ‘mse’, ‘percentile’]
- 关键字参数
stride – 整数。默认为 1
start_bin – 整数。默认为 128
percentils – [0, 100] 之间的浮点数。默认为 99.99。
- Returns
amax – 张量
- reset()
重置收集的直方图
pytorch_quantization.nn
TensorQuantizer
- class pytorch_quantization.nn.TensorQuantizer(quant_desc=<pytorch_quantization.tensor_quant.ScaledQuantDescriptor object>, disabled=False, if_quant=True, if_clip=False, if_calib=False)
张量量化器模块
此模块使用 tensor_quant 或 fake_tensor_quant 函数来量化张量。并包装变量、移动统计信息，我们希望在训练量化网络时使用它们。
- 实验性功能
clip阶段在启用量化之前学习范围。calib阶段运行校准
- 参数
quant_desc –
QuantDescriptor的实例。disabled – 布尔值。如果为 True，则绕过整个模块并返回输入。默认为 False。
if_quant – 布尔值。如果为 True，则运行主量化体。默认为 True。
if_clip – 布尔值。如果为 True，则在量化之前进行裁剪并学习 amax。默认为 False。
if_calib – 布尔值。如果为 True，则运行校准。尚未实现。校准的设置可能会转到
QuantDescriptor。
Raises
- 只读属性
axis
fake_quant
scale
step_size
- 可变属性
num_bits
unsigned
amax
- __init__(quant_desc=<pytorch_quantization.tensor_quant.ScaledQuantDescriptor object>, disabled=False, if_quant=True, if_clip=False, if_calib=False)
初始化量化器并设置所需的变量
- disable()
绕过该模块
- disable_clip()
禁用 clip 阶段
- enable_clip()
启用 clip 阶段
- forward(inputs)
将 tensor_quant 函数应用于输入
- 参数
inputs – 类型为 float32 的张量。
- Returns
outputs – 类型为 output_dtype 的张量
- init_learn_amax()
从固定的 amax 初始化学习到的 amax
量化模块
_QuantConvNd
- class pytorch_quantization.nn.modules.quant_conv._QuantConvNd(in_channels, out_channels, kernel_size, stride, padding, dilation, transposed, output_padding, groups, bias, padding_mode, quant_desc_input, quant_desc_weight)
量化 Conv 的基类，继承自 _ConvNd
原始参数的注释可以在 torch.nn.modules.conv 中找到
- 参数
quant_desc_input –
QuantDescriptor的实例。输入的量化描述符。quant_desc_weight –
QuantDescriptor的实例。权重的量化描述符。
- Raises
ValueError – 如果传入了不支持的参数。
- 只读属性
input_quantizer
weight_quantizer
- 静态方法
set_default_quant_desc_input: 设置 default_quant_desc_input
set_default_quant_desc_weight: 设置 default_quant_desc_weight
QuantConv1d
- class pytorch_quantization.nn.QuantConv1d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', **kwargs)
量化的 1D Conv
QuantConv2d
- class pytorch_quantization.nn.QuantConv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', **kwargs)
量化的 2D conv
QuantConv3d
- class pytorch_quantization.nn.QuantConv3d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', **kwargs)
量化的 3D Conv
QuantConvTranspose1d
- class pytorch_quantization.nn.QuantConvTranspose1d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros', **kwargs)
量化的 ConvTranspose1d
QuantConvTranspose2d
- class pytorch_quantization.nn.QuantConvTranspose2d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros', **kwargs)
量化的 ConvTranspose2d
QuantConvTranspose3d
- class pytorch_quantization.nn.QuantConvTranspose3d(in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, groups=1, bias=True, dilation=1, padding_mode='zeros', **kwargs)
量化的 ConvTranspose3d
QuantLinear
- class pytorch_quantization.nn.QuantLinear(in_features, out_features, bias=True, **kwargs)
nn.Linear 的量化版本
将量化线性应用于传入数据，y = dequant(quant(x)quant(A)^T + b)。
保留模块名称 “Linear” 而不是 “QuantLinear”，以便它可以轻松放入预先存在的模型并加载预训练的权重。别名 “QuantLinear” 在下面定义。基本代码是 nn.Linear 的副本，请参阅原始参数的详细注释。
量化描述符在 kwargs 中传入。如果不存在，则使用 default_quant_desc_input 和 default_quant_desc_weight。
- 关键字参数
quant_desc_input –
QuantDescriptor的实例。输入的量化描述符。quant_desc_wegiht –
QuantDescriptor的实例。权重的量化描述符。
- Raises
ValueError – 如果传入了不支持的参数。
KeyError – 如果传入了不支持的 kwargs。
- 只读属性
input_quantizer
weight_quantizer
- 静态方法
set_default_quant_desc_input: 设置 default_quant_desc_input
set_default_quant_desc_weight: 设置 default_quant_desc_weight
QuantMaxPool1d
- class pytorch_quantization.nn.QuantMaxPool1d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False, **kwargs)
量化的 1D maxpool
QuantMaxPool2d
- class pytorch_quantization.nn.QuantMaxPool2d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False, **kwargs)
量化的 2D maxpool
QuantMaxPool3d
- class pytorch_quantization.nn.QuantMaxPool3d(kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False, **kwargs)
量化的 3D maxpool
QuantAvgPool1d
- class pytorch_quantization.nn.QuantAvgPool1d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, **kwargs)
量化的 1D 平均池化
QuantAvgPool2d
- class pytorch_quantization.nn.QuantAvgPool2d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None, **kwargs)
量化的 2D 平均池化
QuantAvgPool3d
- class pytorch_quantization.nn.QuantAvgPool3d(kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None, **kwargs)
量化的 3D 平均池化
QuantAdaptiveAvgPool1d
- class pytorch_quantization.nn.QuantAdaptiveAvgPool1d(output_size, **kwargs)
量化的 1D 自适应平均池化
QuantAdaptiveAvgPool2d
- class pytorch_quantization.nn.QuantAdaptiveAvgPool2d(output_size, **kwargs)
量化的 2D 自适应平均池化
QuantAdaptiveAvgPool3d
- class pytorch_quantization.nn.QuantAdaptiveAvgPool3d(output_size, **kwargs)
量化的 3D 自适应平均池化
Clip
- class pytorch_quantization.nn.Clip(clip_value_min, clip_value_max, learn_min=False, learn_max=False)
Clip 张量
- 参数
clip_value_min – 用于 clip 的下限的数字或张量
clip_value_max – 用于 clip 的上限的数字或张量
learn_min – 一个布尔值。如果为 True，则学习最小值。clip_value_min 将用于初始化。默认为 False
learn_max – 一个布尔值。与 learn_min 类似，但用于最大值。
- Raises
pytorch_quantization.nn.functional
一些支持函数
ClipFunction
- class pytorch_quantization.nn.functional.ClipFunction(*args, **kwargs)
通用的张量 clip 函数
Pytorch 的 clamp() 仅支持标量范围，且不支持广播。此实现使用 min/max，它更通用。梯度根据 IBM 的 PACT 论文 https://arxiv.org/abs/1805.06085 定义，这也是 Tensorflow 的 clip_by_value() 的行为
clip 是 ClipFunction.apply 的别名
pytorch_quantization.optim.helper
量化优化器/训练器的辅助函数
- pytorch_quantization.optim.helper.freeze_parameters(model, patterns)
如果 patterns 匹配名称，则将 requires_grad 设置为 False
- 参数
model – 一个模块
patterns – 将用于匹配参数名称的字符串列表。如果参数名称包含任何 pattern，它将被冻结。
- pytorch_quantization.optim.helper.group_parameters(model, patterns_list, lrs=None, momentums=None, weight_decays=None)
Group parameters for using per-parameters option in optimizer
Returns a list of dict that matches Pytorch optimizer fashion, see https://pytorch.ac.cn/docs/stable/optim.html#per-parameter-options for more details.
示例
>>> [ >>> {'params': model.base.parameters()}, >>> {'params': model.classifier.parameters(), 'lr': 1e-3} >>> ]
Parameters will be grouped w.r.t first level of the keys_list. e.g. keys_list=[[‘conv1’, ‘conv2’], [‘conv3’]] will return 2 groups, one with conv1 and conv2 in name, and the other with conv3 in name.
If lr, momentum or weight_decay are supplied, they will be added to the group as well.
- 参数
model – A module
patterns_list – A list of list of strings. WARNING: patters must be EXCLUSIVE, the function doesn’t perform exclusive check.
lrs – A list of float with same length as keys_list or None.
momentums – A list of float with same length as keys_list or None.
weight_decays – A list of float with same length as keys_list or None.
- Returns
param_group – A list of dict
- pytorch_quantization.optim.helper.match_parameters(model, patterns)
Returns an generator over module parameters if name matches key
It is useful to group parameters, and apply different functions to different group. This function provides an easy way to group them.
- 参数
model – 一个模块
patterns – A list of strings that will be used to match parameter names. If parameter name contains any pattern, it will be yield
- Yields
param – Module parameters
- pytorch_quantization.optim.helper.quant_weight_inplace(model)
Make quantization inplace
Search for quantized modules including QuantConvNd and QuantLinear, make weight quantization in place using weight_quantizer.
Most publications of quantization aware training uses STE by default, which is really an approximation of derivative of the nondifferentiable quantization function, which works to some extended but by no means the F=ma of the problem. Inplace quantization can be used to implement relax-and-round, which is a common method in Discrete Optimization’s or Integer Programming.
pytorch_quantization.tensor_quant
Basic tensor quantization functions
QuantDescriptor
- pytorch_quantization.tensor_quant.QuantDescriptor
alias of
pytorch_quantization.tensor_quant.ScaledQuantDescriptor
ScaledQuantDescriptor
- class pytorch_quantization.tensor_quant.ScaledQuantDescriptor(num_bits=8, name=None, **kwargs)
Supportive descriptor of quantization
Describe how a tensor should be quantized. A QuantDescriptor and a tensor defines a quantized tensor.
- 参数
num_bits – An integer or a tuple of two integers. Specifically, num_bits can be
- A positive integer argument for integer qunatization. num_bits specify
the number of bits used for integer quantization.
- Constant integer tuple (4,3) for E4M3 floating point quantization emulating
Nvidia’s FP8 quantization. E4M3 quantization only supports per-tensor quantization.
Default: 8.
name – Seems a nice thing to have
- 关键字参数
fake_quant – A boolean. If True, use fake quantization mode. Default True.
axis – None, int or tuple of int. axes which will have its own max for computing scaling factor. If None (the default), use per tensor scale. Must be in the range [-rank(input_tensor), rank(input_tensor)). e.g. For a KCRS weight tensor, quant_axis=(0) will yield per channel scaling. Default None.
amax – A float or list/ndarray of floats of user specified absolute max range. If supplied, ignore quant_axis and use this to quantize. If learn_amax is True, will be used to initialize learnable amax. Default None.
learn_amax – A boolean. If True, learn amax. Default False.
scale_amax – A float. If supplied, multiply amax by scale_amax. Default None. It is useful for some quick experiment.
calib_method – A string. One of [“max”, “histogram”] indicates which calibration to use. Except the simple max calibration, other methods are all hisogram based. Default “max”.
unsigned – A Boolean. If True, use unsigned. Default False.
- Raises
TypeError – If unsupported type is passed in.
- Read-only properties
fake_quant
name
learn_amax
scale_amax
axis
calib_method
num_bits
amax
unsigned
- dict()
Serialize to dict
The build-in __dict__ method returns all the attributes, which includes those have default value and have protected prefix “_”. This method only returns those have values other than the default one and don’t have _ in key. Construct a instance by dict returned by this method should get exactly the same instance.
- classmethod from_yaml(yaml_str)
Create descriptor from yaml str
- to_yaml()
Create yaml serialization Some attributes need special treatment to have human readable form, including amax, axis.
TensorQuantFunction
- class pytorch_quantization.tensor_quant.TensorQuantFunction(*args, **kwargs)
A universal tensor quantization function
Take an input tensor, output an quantized tensor. The granularity of scale can be interpreted from the shape of amax. output_dtype indicates whether the quantized value will be stored in integer or float. The reason we want to store it in float is the pytorch function takes the quantized value may not accept integer input, e.g. Conv2D.
It uses 2^num_bits -1 values instead of 2^num_bits. e.g., for num_bits=8, it uses [-127, 127] instead of [-128, 127]
- static backward(ctx, grad_outputs, grad_scale)
Implements straight through estimation with clipping. For -amax <= input <= amax the gradient passes straight through, otherwise the gradient is zero.
- 参数
ctx – A Context object with saved tensors from forward.
grad_outputs – A tensor of gradient of outputs.
grad_scale – A tensor of gradient of scale.
- Returns
grad_inputs – A tensor of gradient.
- static forward(ctx, inputs, amax, num_bits=8, unsigned=False, narrow_range=True)
Follow tensorflow convention, max value is passed in and used to decide scale, instead of inputing scale directly. Though inputing scale directly may be more natural to use.
- 参数
ctx – A Context object to store tensors for backward.
inputs – 类型为 float32 的张量。
amax – A Tensor of type float32. Inputs will be quantized within range [-amax, amax] amax will be broadcasted to inputs tensor.
num_bits – A integer used to calculate scaling factor, scale = (2^(num_bits-1) - 1) / max Effectively, it indicates how many integer bits is used to represent the value. Default 8.
output_dtype – A type of Tensor. torch.int32 or torch.float32.
unsigned – A boolean. Use unsigned integer range. E.g. [0, 255] for num_bits=8. Default False.
narrow_range – A boolean. Use symmetric integer range for signed quantization E.g. [-127,127] instead of [-128,127] for num_bits=8. Default True.
- Returns
outputs – A Tensor of type output_dtype. scale: A Tensor of type float32. outputs / scale will dequantize outputs tensor.
- Raises
tensor_quant is alias of TensorQuantFunction.apply
fake_tensor_quant is alias of FakeTensorQuantFunction.apply
pytorch_quantization.utils
pytorch_quantization.utils.quant_logging
A WAR for codes that messes up logging format
- pytorch_quantization.utils.quant_logging.reset_logger_handler()
Remove all handler in root logger
pytorch_quantization.utils.reduce_amax
Function to get absolute maximum of a tensor Follow numpy fashion, which is more generic as pytorch’s
- pytorch_quantization.utils.reduce_amax.reduce_amax(input, axis=None, keepdims=True)
Compute the absolute maximum value of a tensor.
Reduces input_tensor along the dimensions given in axis. Unless keepdims is true, the rank of the tensor is reduced by 1 for each entry in axis. If keepdims is true, the reduced dimensions are retained with length 1.
注意
Gradient computeation is disabled as this function is never meant learning reduces amax
- 参数
input – Input tensor
axis – The dimensions to reduce. None or int or tuple of ints. If None (the default), reduces all dimensions. Must be in the range [-rank(input_tensor), rank(input_tensor)).
keepdims – A boolean. If true, retains reduced dimensions with length 1. Default True
granularity – DEPRECTED. specifies if the statistic has to be calculated at tensor or channel granularity
- Returns
The reduced tensor.
- Raises
ValueError – Any axis which doesn’t make sense or is not supported
ValueError – If unknown granularity is passed in.