PyTorch 量化设计方案#

参考：torch_quantization_design_proposal

量化的数学描述#

定义缩放函数（scaling function） \(sc: \mathbb{R}^n \to \mathbf{Q}^n\)，将来自任意范围的值归一化到 \([\mathbf{Q}_{\min}, \mathbf{Q}_{\max}]\)。给定这样的函数，量化函数 (1)。

其中 \(\mathbf{Q}_{\min}\), \(\mathbf{Q}_{\max}\) 表示量化级别的数据类型的最大最小值。比如，对于对称量化，int8 对应 \(-128\)，\(127\)。

量化的表示#

方案：逐张量（per tensor） 非对称线性量化 （asymmetric linear quantization），这意味着张量内的所有值都以相同的方式缩放，输入数据的最小值和最大值线性映射到量化数据类型的最小值和最大值，从而表示零，没有量化误差。

映射是通过

\[ Q(x, \text{scale}, \text{zero_point}) = \operatorname{round}(\cfrac{x}{\text{scale}} + \text{zero_point}) \]

小技巧

\(\text{zero_point}\) 确保浮点数中的 \(0\) 在量化后没有错误表示，从而确保 padding 等运算不会导致额外的量化误差。

变换浮点张量的。

备注

对于算子，限制为：

  1. 8 位权重（data_type = qint8）

  2. 8 位激活（data_type = quint8）

  3. 32 位，对称量化 bias（zero_point = 0, data_type = qint32）

量化张量#

参考：量化张量

为了在 PyTorch 中进行量化，需要能够在张量中表示量化的数据。量化张量允许存储量化数据（表示为 int8/uint8/int32）以及量化参数，如 scale 和 zero_point。量化张量允许许多有用的算子，使得量化算术变得容易，此外还允许以量化格式序列化数据。

关于量化张量的有用函数的简短列表：

  • torch.quantize_linear(x: torch.tensor, scale: float, zero_point: int, dtype: torch.dtype)：使用 scale 和 zero_point 将输入 x 量化为 dtype

  • qx.dequantize()：将张量 qx 反量化为浮点数

  • qx.__repr__()：打印量化张量 qx 的反量化值

  • qx.int_repr()：打印量化张量 qx 中的原始值

创建量化张量#

  • 通过量化非量化的浮点张量得到量化张量

import torch

float_tensor = torch.randn(2, 2, 3)

scale, zero_point = 1e-4, 2
dtype = torch.qint32
q_per_tensor = torch.quantize_per_tensor(float_tensor, scale, zero_point, dtype)
q_per_tensor

tensor([[[ 1.2883, -0.4246, -1.7423],
         [-0.4073,  0.4799, -0.6273]],

        [[-0.1805,  0.8215, -1.5591],
         [-1.7428, -0.6705,  0.0260]]], size=(2, 2, 3), dtype=torch.qint32,
       quantization_scheme=torch.per_tensor_affine, scale=0.0001, zero_point=2)

还支持逐通道量化：

scales = torch.tensor([1e-1, 1e-2, 1e-3])
zero_points = torch.tensor([-1, 0, 1])
channel_axis = 2
q_per_channel = torch.quantize_per_channel(float_tensor,
                                           scales,
                                           zero_points,
                                           axis=channel_axis,
                                           dtype=dtype)
q_per_channel

tensor([[[ 1.3000, -0.4200, -1.7420],
         [-0.4000,  0.4800, -0.6270]],

        [[-0.2000,  0.8200, -1.5590],
         [-1.7000, -0.6700,  0.0260]]], size=(2, 2, 3), dtype=torch.qint32,
       quantization_scheme=torch.per_channel_affine,
       scale=tensor([0.1000, 0.0100, 0.0010], dtype=torch.float64),
       zero_point=tensor([-1,  0,  1]), axis=2)

  • 直接从 empty_quantized 函数创建量化张量

注意，_empty_affine_quantized 是一个私有 API，我们将用类似 torch 的方式替换它。将来使用 empty_quantized_tensor(sizes, quantizer)：

q = torch._empty_affine_quantized([10],
                                  scale=scale,
                                  zero_point=zero_point,
                                  dtype=dtype)
q

tensor([-1.0000e-04, -3.0000e-04, -2.0000e-04,  2.0000e-04,  4.4000e-03,
        -2.0000e-04,  7.9000e-03, -2.0000e-04, -1.3377e+05,  3.2649e+00],
       size=(10,), dtype=torch.qint32,
       quantization_scheme=torch.per_tensor_affine, scale=0.0001, zero_point=2)

  • 通过集合 int 张量和量化参数来创建量化张量

备注

注意，_per_tensor_affine_qtensor 是私有 API，我们将用类似 torch 的东西 torch.form_tensor(int_tensor, quantizer) 替换它

int_tensor = torch.randint(0, 100, size=(10,), dtype=torch.uint8)

数据类型为 torch.quint8，即对应的 torch.uint8，我们有以下对应的 torch int 类型和 torch 量化 int 类型：

  • torch.uint8 -> torch.quint8

  • torch.int8 -> torch.qint8

  • torch.int32 -> torch.qint32

q = torch._make_per_tensor_quantized_tensor(int_tensor, scale, zero_point)  # Note no `dtype`
q 

tensor([0.0055, 0.0067, 0.0039, 0.0042, 0.0077, 0.0003, 0.0036, 0.0007, 0.0090,
        0.0022], size=(10,), dtype=torch.quint8,
       quantization_scheme=torch.per_tensor_affine, scale=0.0001, zero_point=2)

在当前的 API 中，我们必须专一每个量化方案的函数，例如，如果我们想量化张量，我们将有 quantize_per_tensor 和 quantize_per_channel。类似地，对于 q_scale 和 q_zero_point，我们应该有以 Quantizer 作为参数的单一量化函数。为了检查量化参数，我们应该让量化张量返回 Quantizer 对象，这样我们就可以在 Quantizer 对象上检查量化参数，而不是把所有东西都放到张量 API 中。当前的基础设施还没有为这种支持做好准备，目前正在开发中。

量化张量的运算#

反量化

dequantized_tensor = q.dequantize()
dequantized_tensor

tensor([0.0055, 0.0067, 0.0039, 0.0042, 0.0077, 0.0003, 0.0036, 0.0007, 0.0090,
        0.0022])

支持切片

量化张量像通常的张量一样支持切片：

s = q[2]
s

tensor(0.0039, size=(), dtype=torch.quint8,
       quantization_scheme=torch.per_tensor_affine, scale=0.0001, zero_point=2)

备注

尺度（scale）和零点（zero_point）相同的量化张量，它包含与 q_made_per_tensor[2, :] 相同的原始量化张量的第二行值。

赋值

q[0] = 3.5 # 量化 3.5 并将 int 值存储在量化张量中

拷贝

我们可以从量化张量复制相同大小和 dtype 但不同尺度和零点的张量：

scale1, zero_point1 = 1e-1, 0
scale2, zero_point2 = 1, -1
q1 = torch._empty_affine_quantized([2, 3],
                                   scale=scale1,
                                   zero_point=zero_point1,
                                   dtype=torch.qint8)
q2 = torch._empty_affine_quantized([2, 3],
                                   scale=scale2,
                                   zero_point=zero_point2,
                                   dtype=torch.qint8)
q2.copy_(q1)

tensor([[ 1.6000,  2.6000,  2.4000],
        [ 5.8000, -8.8000,  8.5000]], size=(2, 3), dtype=torch.qint8,
       quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=0)

q1.transpose(0, 1)  # see https://pytorch.org/docs/stable/torch.html#torch.transpose
q1.permute([1, 0])  # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.permute
q1.contiguous()  # Convert to contiguous Tensor

tensor([[ 1.6000,  2.6000,  2.4000],
        [ 5.8000, -8.8000,  8.5000]], size=(2, 3), dtype=torch.qint8,
       quantization_scheme=torch.per_tensor_affine, scale=0.1, zero_point=0)

import tempfile
with tempfile.NamedTemporaryFile() as f:
    torch.save(q2, f)
    f.seek(0)
    q3 = torch.load(f)

检查量化张量#

# Check size of Tensor
q.numel(), q.size()

(10, torch.Size([10]))

# Check whether the tensor is quantized
q.is_quantized

True

# Get the scale of the quantized Tensor, only works for affine quantized tensor
q.q_scale()

0.0001

# Get the zero_point of quantized Tensor
q.q_zero_point()

2

# get the underlying integer representation of the quantized Tensor
# int_repr() returns a Tensor of the corresponding data type of the quantized data type
# e.g.for quint8 Tensor it returns a uint8 Tensor while preserving the MemoryFormat when possible
q.int_repr()

tensor([255,  69,  41,  44,  79,   5,  38,   9,  92,  24], dtype=torch.uint8)

# If a quantized Tensor is a scalar we can print the value:
# item() will dequantize the current tensor and return a Scalar of float
q[0].item()

0

# printing
print(q)

tensor([0.0253, 0.0067, 0.0039, 0.0042, 0.0077, 0.0003, 0.0036, 0.0007, 0.0090,
        0.0022], size=(10,), dtype=torch.quint8,
       quantization_scheme=torch.per_tensor_affine, scale=0.0001, zero_point=2)

# indexing
print(q[0]) # q[0] is a quantized Tensor with one value

tensor(0.0253, size=(), dtype=torch.quint8,
       quantization_scheme=torch.per_tensor_affine, scale=0.0001, zero_point=2)

量化的算子/内核#

量化算子，如量化 QRelu、QAdd、QCat、QLinear、QConv 等。要么使用简单的算子实现，要么在算子符中封装 fbgemm 实现。所有的运算都是在 C10 中注册的，而且现在只在 CPU 中。也有关于如何写量化算子/内核的说明。

量化模型#

还有量化的模块，它们封装了这些内核实现，这些内核实现位于 torch.nn.quantized 命名空间中，将在模型开发中使用。提供实用函数来将 torch.nn.Module 替换为 torch.nn.quantized.Module，但用户也可以自由地直接使用它们。尽量将量化模块的 api 与 torch.nn.Module 中的对应 api 匹配。