如何提取qint8静态量化模型权重以在硬件上部署？

你的问题核心是对PyTorch静态量化模型的权重存储格式和量化推理流程存在误解，以下是具体错误点和修正方案：

1. 权重提取方式错误 #

你直接使用state_dict中显示的浮点数（比如[-0.1180, 0.1180]）计算整数权重，但这些数值是PyTorch自动反量化后的浮点数展示，并非原始的qint8整数权重。正确的整数权重应该直接获取tensor的整数表示：

# 以input_layer为例，从state_dict中取出权重tensor
input_weight_tensor = qmodel.state_dict()['input_layer._packed_params._packed_params'][0]
# 获取qint8的原始整数权重
W1_int = input_weight_tensor.int_repr().numpy()
# 对应的scale和zero_point直接从tensor属性获取
W1_scale = input_weight_tensor.scale.item()
W1_zero_point = input_weight_tensor.zero_point.item()

2. 输入未做量化处理 #

量化模型的输入必须先按照输入层的量化参数转换为qint8，而非直接用原始浮点数。输入层的量化参数在state_dict的input_layer_input_scale_0和input_layer_input_zero_point_0中：

input_scale = qmodel.state_dict()['input_layer_input_scale_0'].item()
input_zero_point = qmodel.state_dict()['input_layer_input_zero_point_0'].item()

# 量化输入：x_q = round(x / input_scale) + input_zero_point，裁剪到qint8范围
inputs_q = np.clip(np.round(inputs / input_scale) + input_zero_point, -128, 127).astype(np.int8)

3. 偏置的处理逻辑错误 #

PyTorch量化线性层的偏置是浮点数值，已经考虑了输入和权重的scale，不需要再除以权重的scale转整数，直接取出使用即可：

# 获取input_layer的偏置（浮点）
b1 = qmodel.state_dict()['input_layer._packed_params._packed_params'][1].numpy()
# 获取out层的偏置（浮点）
b2 = qmodel.state_dict()['out._packed_params._packed_params'][1].numpy()

4. 量化推理的完整流程错误 #

量化模型的推理需要严格遵循量化计算-反量化-量化-计算的流程，用浮点模拟量化运算时，必须正确应用scale的乘积。完整的正确推理步骤如下：

步骤1：输入量化后计算第一层线性输出 #

# 整数矩阵乘法（输入qint8 × 权重qint8）
linear1_int = inputs_q @ W1_int.T
# 反量化为浮点：因权重和输入的zero_point都是0，简化为整数结果 × 输入scale × 权重scale
linear1_float = linear1_int * input_scale * W1_scale
# 加偏置后应用ReLU
hidden_float = np.maximum(linear1_float + b1, 0)

步骤2：隐藏层输出量化后计算第二层线性输出 #

首先获取隐藏层输出的量化参数（对应state_dict中的input_layer.scale和input_layer.zero_point，即第一层输出的量化参数）：

hidden_scale = qmodel.state_dict()['input_layer.scale'].item()
hidden_zero_point = qmodel.state_dict()['input_layer.zero_point'].item()

# 量化隐藏层输出
hidden_q = np.clip(np.round(hidden_float / hidden_scale) + hidden_zero_point, -128, 127).astype(np.int8)

# 获取输出层的整数权重、scale
output_weight_tensor = qmodel.state_dict()['out._packed_params._packed_params'][0]
W2_int = output_weight_tensor.int_repr().numpy()
W2_scale = output_weight_tensor.scale.item()

# 整数矩阵乘法
linear2_int = hidden_q @ W2_int.T
# 反量化为浮点
linear2_float = linear2_int * hidden_scale * W2_scale
# 加偏置得到最终输出
output_float = linear2_float + b2

步骤3：计算分类结果 #

predictions = output_float.argmax(axis=1)
print(predictions)

完整修正后的代码示例 #

import numpy as np
import torch

# 假设已加载量化后的qmodel
# 提取各层量化参数
# 输入层量化参数
input_scale = qmodel.state_dict()['input_layer_input_scale_0'].item()
input_zero_point = qmodel.state_dict()['input_layer_input_zero_point_0'].item()

# 第一层权重与偏置
input_weight_tensor = qmodel.state_dict()['input_layer._packed_params._packed_params'][0]
W1_int = input_weight_tensor.int_repr().numpy()
W1_scale = input_weight_tensor.scale.item()
b1 = qmodel.state_dict()['input_layer._packed_params._packed_params'][1].numpy()

# 第一层输出量化参数（隐藏层输入量化参数）
hidden_scale = qmodel.state_dict()['input_layer.scale'].item()
hidden_zero_point = qmodel.state_dict()['input_layer.zero_point'].item()

# 输出层权重与偏置
output_weight_tensor = qmodel.state_dict()['out._packed_params._packed_params'][0]
W2_int = output_weight_tensor.int_repr().numpy()
W2_scale = output_weight_tensor.scale.item()
b2 = qmodel.state_dict()['out._packed_params._packed_params'][1].numpy()

# 测试输入
inputs = np.array(
       [[1.  , 1.  ], # class 0 example
       [1.  , 0.  ], # class 1 example
       [0.  , 1.  ],
       [0.  , 0.  ],
       [0.  , 0.9 ],
       [0.  , 0.75],
       [0.  , 0.25]]) # class 6 example

# 量化输入
inputs_q = np.clip(np.round(inputs / input_scale) + input_zero_point, -128, 127).astype(np.int8)

# 第一层计算
linear1_int = inputs_q @ W1_int.T
linear1_float = linear1_int * input_scale * W1_scale
hidden_float = np.maximum(linear1_float + b1, 0)

# 量化隐藏层输出
hidden_q = np.clip(np.round(hidden_float / hidden_scale) + hidden_zero_point, -128, 127).astype(np.int8)

# 第二层计算
linear2_int = hidden_q @ W2_int.T
linear2_float = linear2_int * hidden_scale * W2_scale
output_float = linear2_float + b2

# 获取预测结果
predictions = output_float.argmax(axis=1)
print(predictions)

额外说明 #

• PyTorch静态量化模型的state_dict中，每个线性层的_packed_params包含量化后的权重tensor（qint8）和浮点偏置，权重tensor的scale和zero_point是该权重的量化参数。
• 输入层的input_layer_input_scale_0是输入的量化参数，input_layer.scale是第一层输出的量化参数，也就是第二层输入的量化参数。
• 输出层的out.scale是最终输出的量化参数，若需浮点输出做argmax，可直接使用反量化后的浮点结果，无需额外处理。

内容的提问来源于stack exchange，提问作者Chrispresso