一、Fundamentals
- Brevitas中使用量化有两种方法，一种是使用量化器，一种是直接设置参数
1.QuantLinear层
brevitas.nn.QuantLinear是torch.nn.Linear的量化变体形式，同时还是QuantWeightBiasInputOutputLayer的实例，意味着其支持权重、偏置、输入和输出的量化。其他的实例还包括QuantConv1d,QuantConv2d,QuantConvTranspose1d,QuantConvTranspose2d,他们都遵循相同的原则。
import inspect
from brevitas.nn import QuantLinear
from IPython.display import Markdown, display
def pretty_print_source(source):
display(Markdown('```python\n' + source + '\n```'))
source = inspect.getsource(QuantLinear.__init__)
pretty_print_source(source)
默认情况下权重量化weight_quant = Int8WeightPerTensorFloat,其余bisa_quant,input_quant,output_quant,默认为None。
即默认情况下只有权重为8位量化，其余不量化，Int8WeightPerTensorFloat是量化器。
2.权重量化
不同的量化器有不同的量化策略。几个例子：Int8WeightPerTensorFloat根据要量化的浮点权重张量中找到的最大值来计算比例因子(scale)；Int8WeightPerTensorFixedPoint限制scale为2的幂次方；SignedBinaryWeightPerTensorConst的scale为固定的0.1.
2.1 默认权重量化
当我们不对网络结构使用任何量化器时，权重会执行默认量化：weight_quant = Int8WeightPerTensorFloat，而偏置、输入和输出量化则被禁用。
import torch
torch.manual_seed(0)
quant_linear = QuantLinear(2, 4, bias=True)
print(f"Original float weight tensor:\n {quant_linear.weight} \n")
print(f"Quantized weight QuantTensor:\n {quant_linear.quant_weight()} \n")
Original float weight tensor:
Parameter containing:
tensor([[-0.0053, 0.3793],
[-0.5820, -0.5204],
[-0.2723, 0.1896],
[-0.0140, 0.5607]], requires_grad=True)
Quantized weight QuantTensor:
QuantTensor(value=tensor([[-0.0046, 0.3803],
[-0.5820, -0.5224],
[-0.2704, 0.1879],
[-0.0137, 0.5591]], grad_fn=<MulBackward0>), scale=tensor(0.0046, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
QuantTensor相比于一般的张量，包含了量化相关的一些参数。比如：scale是缩放比，用于全精度与量化值之间的转换；zero_point（还没弄明白时干啥的）；bit_width参数位宽；signed_t参数是否为带符号数；training_t(应该是是否可训练的意思)。
2.2 权重量化+浮点输入
仅启用权重量化，传入浮点输入，会得到浮点输出。浮点输出是由浮点输入与反量化权重计算得来。反量化权重：网络层内权重是以量化后的形式存储的，但在推理计算的时候，将存储的权重反量化后参与计算。
2.3 定点量化
Int8WeightPerTensorFixedPoint
2.4 二值量化
量化器为：SignedBinaryWeightPerTensorConst
torch.manual_seed(0)
from brevitas.quant import SignedBinaryWeightPerTensorConst
quant_linear = QuantLinear(2, 4, weight_quant=SignedBinaryWeightPerTensorConst, bias=False)
print(f"Weight QuantTensor:\n {quant_linear.quant_weight()}")
Weight QuantTensor:
QuantTensor(value=tensor([[-0.1000, 0.1000],
[-0.1000, -0.1000],
[-0.1000, 0.1000],
[-0.1000, 0.1000]], grad_fn=<MulBackward0>), scale=tensor(0.1000), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
量化后由固定的缩放比scale = 0.1，位宽bit_width = 1。
2.5 权重量化器的共享
Brevitas 还允许在层之间共享权重量化器的实例（而不是定义），它迫使它们具有相同的尺度、零点和位宽。在多个层之间共享它意味着量化器现在查看所有正在量化的权重张量以确定总体最大值并生成单个比例因子。（共享就是使用实例过的层的量化器，如quant_linear1.weight_quant）
torch.manual_seed(0)
# Define a QuantLinear layer 1
quant_linear1 = QuantLinear(2, 4, bias=False)
# Keep a pointer to the scale factor of QuantLinear layer 1 weights before sharing
quant_linear1_scale_before_sharing = quant_linear1.quant_weight().scale
# Define a QuantLinear layer 2 where the weight quantizer is taken from layer 1
quant_linear2 = QuantLinear(2, 4, weight_quant=quant_linear1.weight_quant, bias=False)
print(f"QuantLinear 1 scale before sharing with QuantLinear 2: {quant_linear1_scale_before_sharing:.4f}")
print(f"QuantLinear 2 scale: {quant_linear2.quant_weight().scale:.4f}")
print(f"QuantLinear 1 scale after sharing with QuantLinear 2: {quant_linear1.quant_weight().scale:.4f}")
QuantLinear 1 scale before sharing with QuantLinear 2: 0.0046
QuantLinear 2 scale: 0.0053
QuantLinear 1 scale after sharing with QuantLinear 2: 0.0053
2.输入、输出激活量化
生成量化输入，则会生成量化输出。可以设置输入量化器，如:Int8ActPerTensorFloat
torch.manual_seed(0)
from brevitas.quant import Int8ActPerTensorFloat
float_input = torch.randn(3, 2)
quant_linear = QuantLinear(2, 4, input_quant=Int8ActPerTensorFloat, bias=False)
quant_output = quant_linear(float_input)
print(f"Float input:\n {float_input} \n")
print(f"Quant output:\n {quant_output}")
Float input:
tensor([[ 1.5410, -0.2934],
[-2.1788, 0.5684],
[-1.0845, -1.3986]])
Quant output:
tensor([[-0.9109, -0.4609, 0.3135, -0.6523],
[ 1.2089, 0.6524, -0.3752, 0.8697],
[ 1.3893, 0.2816, -0.9011, 0.9521]], grad_fn=<MmBackward>)
可以观察到输出是个普通的tensor，并没有包含量化输出的信息。可以设置return_quant_tensor=True属性就会返回一个QuantTensor。
torch.manual_seed(0)
from brevitas.quant import Int8ActPerTensorFloat
float_input = torch.randn(3, 2)
quant_linear = QuantLinear(2, 4, input_quant=Int8ActPerTensorFloat, bias=False, return_quant_tensor=True)
quant_output = quant_linear(float_input)
print(f"Quant output:\n {quant_output}")
Quant output:
QuantTensor(value=tensor([[-0.9109, -0.4609, 0.3135, -0.6523],
[ 1.2089, 0.6524, -0.3752, 0.8697],
[ 1.3893, 0.2816, -0.9011, 0.9521]], grad_fn=<MmBackward>), scale=tensor([[9.0542e-05]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17.), signed_t=tensor(True), training_t=tensor(True))
2.1 QuantIdentity 层
对于输入量化，也可以在输入之前加上一个QuantIdentity层。QuantIdentity层默认情况下使用量化器Int8ActPerTensorFloat对输出进行量化。
torch.manual_seed(0)
from brevitas.nn import QuantIdentity
float_input = torch.randn(3, 2)
quant_identity = QuantIdentity(return_quant_tensor=True)
quant_linear = QuantLinear(2, 4, bias=False, return_quant_tensor=True)
quant_input = quant_identity(float_input)
quant_output = quant_linear(quant_input)
print(f"Float input:\n {float_input} \n")
print(f"Quant input:\n {quant_input} \n")
Float input:
tensor([[ 1.5410, -0.2934],
[-2.1788, 0.5684],
[-1.0845, -1.3986]])
Quant input:
QuantTensor(value=tensor([[ 1.5490, -0.2894],
[-2.1788, 0.5617],
[-1.0894, -1.3958]], grad_fn=<MulBackward0>), scale=tensor(0.0170, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
Quant output:
QuantTensor(value=tensor([[-0.9109, -0.4609, 0.3135, -0.6523],
[ 1.2089, 0.6524, -0.3752, 0.8697],
[ 1.3893, 0.2816, -0.9011, 0.9521]], grad_fn=<MmBackward>), scale=tensor([[9.0542e-05]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17.), signed_t=tensor(True), training_t=tensor(True))
2.2 QuantReLU层
QuantReLU 拥有默认的输出量化Uint8ActPerTensorFloat，意味着QuantReLU 先进行relu计算，然后对结果进行量化并输出。
torch.manual_seed(0)
from brevitas.nn import QuantReLU
float_input = torch.randn(3, 2)
quant_relu = QuantReLU(return_quant_tensor=True)
quant_output = quant_relu(float_input)
print(f"Float input:\n {float_input} \n")
print(f"Quant output:\n {quant_output}")
Float input:
tensor([[ 1.5410, -0.2934],
[-2.1788, 0.5684],
[-1.0845, -1.3986]])
Quant output:
QuantTensor(value=tensor([[1.5410, 0.0000],
[0.0000, 0.5681],
[0.0000, 0.0000]], grad_fn=<MulBackward0>), scale=tensor(0.0060, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(False), training_t=tensor(True))
2.3 默认激活量化的scale如何确定
默认量化器（Uint8ActPerTensorFloat Int8ActPerTensorFloat）通过收集多个训练步骤的统计数据（默认为 300 个步骤的绝对值的 99.999 百分位）来初始化激活的scale.
3.偏置量化
在许多推理工具链中，偏置量化依赖于量化输入和量化权重的scale，这意味着进行偏置量化必须有一个量化输入。brevitas.quant.scaled_int.Int16Bias是一个预定义的偏置量化器，如果只设置了bias_quant=Int16Bias，而没有对输入进行量化，则会报错：
torch.manual_seed(0)
from brevitas.quant.scaled_int import Int16Bias
float_input = torch.randn(3, 2)
quant_linear = QuantLinear(2, 4, bias=True, bias_quant=Int16Bias, return_quant_tensor=True)
quant_output = quant_linear(float_input)
可以通过传入一个QuantTensor或者设置input_quant解决以上报错：
torch.manual_seed(0)
float_input = torch.randn(3, 2)
quant_linear = QuantLinear(
2, 4, bias=True, input_quant=Int8ActPerTensorFloat, bias_quant=Int16Bias, return_quant_tensor=True)
quant_linear(float_input)
QuantTensor(value=tensor([[-0.6541, 0.1263, 0.1680, -0.1231],
[ 1.4658, 1.2395, -0.5207, 1.3989],
[ 1.6461, 0.8687, -1.0466, 1.4813]], grad_fn=<AddmmBackward>), scale=tensor([[9.0542e-05]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(18.), signed_t=tensor(True), training_t=tensor(True))
4.对QuantTensor的操作
4.1 逐元素操作
对于逐元素加法，与传统的定点运算一致，要求操作数的小数位数相同。
torch.manual_seed(0)
float_inp1 = torch.randn(3, 2)
float_inp2 = torch.randn(3, 2)
quant_identity = QuantIdentity(return_quant_tensor=True)
#Training mode, statistics are being collected, scaling factors are different but it doesn't raise an error
train_quant_inp1 = quant_identity(float_inp1)
train_quant_inp2 = quant_identity(float_inp2)
train_mode_add = train_quant_inp1 + train_quant_inp2
#Inference mode, the EMA buffer is being used, scaling factors are the same
quant_identity.eval()
eval_quant_inp1 = quant_identity(float_inp1)
eval_quant_inp2 = quant_identity(float_inp2)
eval_mode_add = eval_quant_inp1 + eval_quant_inp2
print(f"Eval mode add quant inputs:\n {eval_quant_inp1} \n {eval_quant_inp2} \n")
print(f"Eval mode add quant output:\n {eval_mode_add}")
Eval mode add quant inputs:
QuantTensor(value=tensor([[ 1.5335, -0.2875],
[-2.0447, 0.5751],
[-1.0863, -1.4057]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False))
QuantTensor(value=tensor([[ 0.3994, 0.8307],
[-0.7188, -0.3994],
[-0.5910, 0.1757]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False))
Eval mode add quant output:
QuantTensor(value=tensor([[ 1.9329, 0.5431],
[-2.7636, 0.1757],
[-1.6773, -1.2300]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(9.), signed_t=tensor(True), training_t=tensor(False))
4.2 QuantTensor调用torch函数
通过 torch_function 接口（PyTorch >= 1.5.0 支持），可以在 QuantTensor上调用标准 torch 函数。目前，对于仿射量化不变的受支持操作，将返回 QuantTensor，否则输出将返回浮点 torch.Tensor。
max_pool
例如，torch.nn.function.max_pool1d 对于量化是不变的，输出返回QuantTensor:
torch.manual_seed(0)
float_inp = torch.randn(3, 2, 4)
quant_identity = QuantIdentity(return_quant_tensor=True)
quant_input = quant_identity(float_inp)
quant_output = torch.nn.functional.max_pool1d(quant_input, kernel_size=2, stride=2)
print(f"Quant input:\n {quant_input} \n")
print(f"Quant output:\n {quant_output}")
Quant input:
QuantTensor(value=tensor([[[-1.1218, -1.1580, -0.2533, -0.4343],
[ 0.8504, 0.6876, -0.3076, -2.1170]],
[[ 0.4704, -0.1628, 1.4475, 0.2714],
[ 0.1628, 0.8685, -0.1448, -0.1086]],
[[ 0.9228, 1.2666, 2.0084, 0.0543],
[ 0.6152, -0.4162, -0.8323, -2.3160]]], grad_fn=<MulBackward0>), scale=tensor(0.0181, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
Quant output:
QuantTensor(value=tensor([[[-1.1218, -0.2533],
[ 0.8504, -0.3076]],
[[ 0.4704, 1.4475],
[ 0.8685, -0.1086]],
[[ 1.2666, 2.0084],
[ 0.6152, -0.8323]]], grad_fn=<SqueezeBackward1>), scale=tensor(0.0181, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
tanh
但是torch.tanh则不同。他对量化改变，因此输出结果是浮点张量：
torch.manual_seed(0)
float_inp = torch.randn(3, 2, 4)
quant_identity = QuantIdentity(return_quant_tensor=True)
quant_input = quant_identity(float_inp)
quant_output = torch.tanh(quant_input)
print(f"Quant input:\n {quant_input} \n")
print(f"Quant output:\n {quant_output}")
Quant input:
QuantTensor(value=tensor([[[-1.1218, -1.1580, -0.2533, -0.4343],
[ 0.8504, 0.6876, -0.3076, -2.1170]],
[[ 0.4704, -0.1628, 1.4475, 0.2714],
[ 0.1628, 0.8685, -0.1448, -0.1086]],
[[ 0.9228, 1.2666, 2.0084, 0.0543],
[ 0.6152, -0.4162, -0.8323, -2.3160]]], grad_fn=<MulBackward0>), scale=tensor(0.0181, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
Quant output:
tensor([[[-0.8082, -0.8204, -0.2480, -0.4089],
[ 0.6913, 0.5964, -0.2983, -0.9714]],
[[ 0.4386, -0.1614, 0.8952, 0.2649],
[ 0.1614, 0.7006, -0.1438, -0.1081]],
[[ 0.7272, 0.8529, 0.9646, 0.0542],
[ 0.5478, -0.3937, -0.6817, -0.9807]]], grad_fn=<TanhBackward>)
QuantTensor的拼接
当QuantTensor拥有相同的sign, scale, zero-point和 bit-width，就可以使用torch.cat进行张量拼接。训练模式下scale和zero-point可以不同，但在推理模式下必须相同。
torch.manual_seed(0)
float_inp1 = torch.randn(3, 2)
float_inp2 = torch.randn(3, 2)
quant_identity = QuantIdentity(return_quant_tensor=True)
#Training mode, statistics are being collected, scaling factors are different but it doesn't raise an error
train_mode_cat = torch.cat([quant_identity(float_inp1), quant_identity(float_inp2)], dim=1)
#Inference mode, the EMA buffer is being used, scaling factors are the same
quant_identity.eval()
eval_quant_inp1 = quant_identity(float_inp1)
eval_quant_inp2 = quant_identity(float_inp2)
eval_mode_cat = torch.cat([eval_quant_inp1, eval_quant_inp2], dim=1)
print(f"Eval mode concat quant inputs:\n {eval_quant_inp1} {eval_quant_inp2} \n")
print(f"Eval mode concat quant output:\n {eval_mode_cat}")
Eval mode concat quant inputs:
QuantTensor(value=tensor([[ 1.5335, -0.2875],
[-2.0447, 0.5751],
[-1.0863, -1.4057]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False)) QuantTensor(value=tensor([[ 0.3994, 0.8307],
[-0.7188, -0.3994],
[-0.5910, 0.1757]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False))
Eval mode concat quant output:
QuantTensor(value=tensor([[ 1.5335, -0.2875, 0.3994, 0.8307],
[-2.0447, 0.5751, -0.7188, -0.3994],
[-1.0863, -1.4057, -0.5910, 0.1757]]), scale=tensor(0.0160), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(False))
5.自定义量化器
自定义量化器的最简单方法是传递适当的关键字参数。其实就是我们传入参数，覆盖掉底层的量化器的参数。
权重位宽
要覆盖掉weight_quant设置的权重量化器的位宽，需要设置weight_bit_width参数。
torch.manual_seed(0)
quant_linear = QuantLinear(2, 4, weight_bit_width=5, bias=True)
print(f"Weight QuantTensor:\n {quant_linear.quant_weight()}")
权重 QuantTensor:
QuantTensor(value=tensor([[-0.0000, 0.3880],
[-0.5820, -0.5044],
[-0.2716, 0.1940],
[-0.0000, 0.5432]], grad_fn=<MulBackward0>), scale=tensor (0.0388，grad_fn=<DivBackward0>)，zero_point=张量(0.)，bit_width=张量(5.)，signed_t=张量(True)，training_t=张量(True))
上述例子设置weight_bit_width=5，覆盖掉了默认量化器Int8WeightPerTensorFloat的8位权重位宽，输出的结果也验证了这一点。
Per-channel weight quantization
设置weight_scaling_per_output_channel=True启用通道量化：
torch.manual_seed(0)
quant_linear = QuantLinear(2, 4, weight_bit_width=5, weight_scaling_per_output_channel=True, bias=False)
print(f"Weight QuantTensor:\n {quant_linear.quant_weight()}")
Weight QuantTensor:
QuantTensor(value=tensor([[-0.0000, 0.3793],
[-0.5820, -0.5044],
[-0.2723, 0.1816],
[-0.0000, 0.5607]], grad_fn=<MulBackward0>), scale=tensor([[0.0253],
[0.0388],
[0.0182],
[0.0374]], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(5.), signed_t=tensor(True), training_t=tensor(True))
激活位宽
设置bit_width参数改变激活位宽：
torch.manual_seed(0)
float_input = torch.randn(3, 2)
quant_identity = QuantIdentity(bit_width=3, return_quant_tensor=True)
print(f"QuantTensor:\n {quant_identity(float_input)}")
QuantTensor:
QuantTensor(value=tensor([[ 1.6341, -0.5447],
[-2.1788, 0.5447],
[-1.0894, -1.6341]], grad_fn=<MulBackward0>), scale=tensor(0.5447, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(3.), signed_t=tensor(True), training_t=tensor(True))
传入初始化参数max_val进行激活量化
某些量化器要求用户传递额外的关键字参数，比如Uint8ActPerTensorFloatMaxInit，结果就是scale由用户定义的参数max_val决定而不是来自有统计。
torch.manual_seed(0)
from brevitas.quant import Uint8ActPerTensorFloatMaxInit
float_inp1 = torch.randn(3, 2)
quant_relu = QuantReLU(max_val=6.0, act_quant=Uint8ActPerTensorFloatMaxInit, return_quant_tensor=True)
quant_relu(float_inp1)
QuantTensor(value=tensor([[1.5294, 0.0000],
[0.0000, 0.5647],
[0.0000, 0.0000]], grad_fn=<MulBackward0>), scale=tensor(0.0235, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(False), training_t=tensor(True))
Per-channel激活量化
略，以后再补
量化器的继承
可以通过继承我们正在定制的量化器来简单地定义一个新的量化器：
torch.manual_seed(0)
from brevitas.nn import QuantConv1d
BATCHES = 1
CHANNELS = 2
FEATURES = 5
KERNEL = 3
class PerChannel3bActQuant(Int8ActPerTensorFloat):
bit_width = 3
scaling_per_output_channel=True
scaling_stats_permute_dims=(1, 0, 2)
float_input = torch.randn(BATCHES, CHANNELS, FEATURES)
per_channel_depthwise_quant_conv = QuantConv1d(
CHANNELS, CHANNELS, KERNEL, groups=CHANNELS, bias=True,
# set the quantizers
input_quant=PerChannel3bActQuant,
bias_quant=Int16Bias,
# layer-specific kwarg
input_per_channel_broadcastable_shape=(1, CHANNELS, 1),
return_quant_tensor=True)
quant_output = per_channel_depthwise_quant_conv(float_input)
print(f"Float input:\n {float_input} \n")
print(f"Per-channel quant output:\n {quant_output}")
Float input:
tensor([[[ 1.5410, -0.2934, -2.1788, 0.5684, -1.0845],
[-1.3986, 0.4033, 0.8380, -0.7193, -0.4033]]])
Per-channel quant output:
QuantTensor(value=tensor([[[ 0.8616, -0.7012, 0.4503],
[-1.1285, -0.4937, -0.1901]]], grad_fn=<SqueezeBackward1>), scale=tensor([[[0.0021],
[0.0013]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17.), signed_t=tensor(True), training_t=tensor(True))
6.使用枚举从头开始定义量化器
这部分主要是介绍量化器具体是如何定义的，当然这不是唯一的方法，但是却是用户轻松尝试不同内置选项的一种方法。
权重量化器
对于权重，可以通过继承WeightQuantSolver来实现定义量化器：
from brevitas.inject.enum import *
from brevitas.core.zero_point import ZeroZeroPoint
from brevitas.quant.solver import WeightQuantSolver, ActQuantSolver
class Int8WeightPerTensorFloat(WeightQuantSolver):
quant_type = QuantType.INT # integer quantization
bit_width_impl_type = BitWidthImplType.CONST # constant bit width
float_to_int_impl_type = FloatToIntImplType.ROUND # round to nearest
scaling_impl_type = ScalingImplType.STATS # scale based on statistics
scaling_stats_op = StatsOp.MAX # scale statistics is the absmax value
restrict_scaling_type = RestrictValueType.FP # scale factor is a floating point value
scaling_per_output_channel = False # scale is per tensor
bit_width = 8 # bit width is 8
signed = True # quantization range is signed
narrow_range = True # quantization range is [-127,127] rather than [-128, 127]
zero_point_impl = ZeroZeroPoint # zero point is 0.
激活量化器
激活可以通过继承ActQuantSolver来实现定义量化器：
class Int8ActPerTensorFloat(ActQuantSolver):
quant_type = QuantType.INT # integer quantization
bit_width_impl_type = BitWidthImplType.CONST # constant bit width
float_to_int_impl_type = FloatToIntImplType.ROUND # round to nearest
scaling_impl_type = ScalingImplType.PARAMETER_FROM_STATS # scale is a parameter initialized from statistics
scaling_stats_op = StatsOp.PERCENTILE # scale statistics is a percentile of the abs value
high_percentile_q = 99.999 # percentile is 99.999
collect_stats_steps = 300 # statistics are collected for 300 forward steps before switching to a learned parameter
restrict_scaling_type = RestrictValueType.FP # scale is a floating-point value
scaling_per_output_channel = False # scale is per tensor
bit_width = 8 # bit width is 8
signed = True # quantization range is signed
narrow_range = False # quantization range is [-128, 127] rather than [-127, 127]
zero_point_impl = ZeroZeroPoint # zero point is 0.
上述量化器中的任何属性都可以作为关键字参数传入或覆盖量化器要传递到的层（连同其适当的前缀），关键字参数始终优先于量化器中定义的值。
可学习的scale和位宽可变的量化器（example）
per-channel scale factors learned in log domain as a parameter initialized from absmax statistics. - bit-width initialized to 8b and learned as a parameter from there.
torch.manual_seed(0)
from brevitas.quant import Int8WeightPerTensorFloat
class LearnedIntWeightPerChannelFloat(Int8WeightPerTensorFloat):
scaling_per_output_channel = True
scaling_impl_type = ScalingImplType.PARAMETER_FROM_STATS
restrict_scaling_type = RestrictValueType.LOG_FP
bit_width_impl_type = BitWidthImplType.PARAMETER
quant_linear = QuantLinear(2, 4, weight_quant=LearnedIntWeightPerChannelFloat, bias=False)
print(f"Weight QuantTensor:\n {quant_linear.quant_weight()}")
Weight QuantTensor:
QuantTensor(value=tensor([[-0.0060, 0.3793],
[-0.5820, -0.5224],
[-0.2723, 0.1887],
[-0.0132, 0.5607]], grad_fn=<MulBackward0>), scale=tensor([[0.0030],
[0.0046],
[0.0021],
[0.0044]], grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8., grad_fn=<RoundSteFnBackward>), signed_t=tensor(True), training_t=tensor(True))
观察到bit_width是带梯度信息的，因此在反向传播的过程中可以加入到损失函数中进行改变。同样，激活也可以这么应用：
torch.manual_seed(0)
class LearnedIntActPerTensorFloat(Int8ActPerTensorFloat):
bit_width_impl_type = BitWidthImplType.PARAMETER
restrict_scaling_type = RestrictValueType.LOG_FP
float_inp = torch.randn(3, 2)
quant_linear = QuantLinear(
2, 4,
input_quant=LearnedIntActPerTensorFloat,
weight_quant=LearnedIntWeightPerChannelFloat,
return_quant_tensor=True, bias=False)
quant_linear(float_inp)
QuantTensor(value=tensor([[-0.9109, -0.4588, 0.3119, -0.6530],
[ 1.2089, 0.6493, -0.3731, 0.8706],
[ 1.3893, 0.2823, -0.8979, 0.9543]], grad_fn=<MmBackward>), scale=tensor([[9.0542e-05, 3.9068e-05, 5.6866e-05, 6.4251e-05]],
grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(17., grad_fn=<CeilSteFnBackward>), signed_t=tensor(True), training_t=tensor(True))
我们可以通过将bit_width加入到损失函数中进行反向传播，进而自动控制位宽。
从浮点重新训练
在许多场景中，从浮点模型开始执行量化感知训练很方便。 举例来说，我们想要在层的顶部加载一个预训练的浮点状态，并使用我们刚刚看到的学习的位宽和比例。 我们用一个单独的浮点nn.Linear来模拟它。 如果我们不做任何其他事情，我们会得到一个错误：
torch.manual_seed(0)
from torch import nn
float_linear = nn.Linear(2, 4, bias=False)
quant_linear = QuantLinear(
2, 4,
input_quant=LearnedIntActPerTensorFloat,
weight_quant=LearnedIntWeightPerChannelFloat,
return_quant_tensor=True, bias=False)
quant_linear.load_state_dict(float_linear.state_dict())
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
C:\Users\ALESSA~1\AppData\Local\Temp/ipykernel_18920/1653109852.py in <module>
10 return_quant_tensor=True, bias=False)
11
---> 12 quant_linear.load_state_dict(float_linear.state_dict())
~\miniconda3\envs\pt190\lib\site-packages\torch\nn\modules\module.py in load_state_dict(self, state_dict, strict)
1405 if len(error_msgs) > 0:
1406 raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-> 1407 self.__class__.__name__, "\n\t".join(error_msgs)))
1408 return _IncompatibleKeys(missing_keys, unexpected_keys)
1409
RuntimeError: Error(s) in loading state_dict for QuantLinear:
Missing key(s) in state_dict: "input_quant.fused_activation_quant_proxy.tensor_quant.scaling_impl.value", "input_quant.fused_activation_quant_proxy.tensor_quant.msb_clamp_bit_width_impl.bit_width_offset", "weight_quant.tensor_quant.scaling_impl.value", "weight_quant.tensor_quant.msb_clamp_bit_width_impl.bit_width_offset".
这是因为权重和输入的量化器引入了原始浮点层中不存在的新学习参数。
解决办法：1.设置环境变量BREVITAS_IGNORE_MISSING_KEYS=1; 2.通过启用相应的配置标志config.IGNORE_MISSING_KEYS = True.
torch.manual_seed(0)
from torch import nn
from brevitas import config
config.IGNORE_MISSING_KEYS = True
float_linear = nn.Linear(2, 4, bias=False)
quant_linear = QuantLinear(
2, 4,
input_quant=LearnedIntActPerTensorFloat,
weight_quant=LearnedIntWeightPerChannelFloat,
return_quant_tensor=True, bias=False)
quant_linear.load_state_dict(float_linear.state_dict())
7.基于新的量化组件表达完全新颖的算法
不想看了，我研究所需要的部分看完了，其他的有机会在看
8.导出
在使用FINN推理框架之前，需要将网络模型转为Brevitas训练模型，并导出ONNX格式，因为FINN所有的操作都是在ONNX格式的模型下进行的，
Brevitas 通过利用 PyTorch 对自定义符号表示（特别是 ONNX）的支持，支持不同抽象级别的不同风格的导出。
现有的导出流假设静态量化，即比例、零点和位宽需要独立于输入。所有导出流都从如何确定标度、零点和位宽的细节中抽象出来。然而，不同的导出流仅对模型的尺度、零点、精度或结构的某些组合提供支持。
安装导出所需的 onnx 和 onnxoptimizer，并使用 Netron 可视化导出的模型：
!pip install netron onnx onnxoptimizer
!pip install netron onnx onnxoptimizer
已满足要求：c:\users\alessandro\miniconda3\envs\pt190\lib\site-packages 中的 netron (5.3.9) 已满足要求：
c:\users\alessandro\miniconda3\envs\pt190\lib\ 中的 onnx site-packages (1.10.2)
已满足要求： c:\users\alessandro\miniconda3\envs\pt190\lib\site-packages 中的 onnxoptimizer (0.2.6)
已满足要求： c: 中的 numpy>=1.16.6: \users\alessandro\miniconda3\envs\pt190\lib\site-packages (来自 onnx) (1.21.2)
已满足要求：c:\users\alessandro\miniconda3\envs\pt190 中的打字扩展 >=3.6.2.1 \lib\site-packages（来自 onnx）（3.10.0.2）
已满足要求：c:\users\alessandro\miniconda3\envs\pt190\lib\site-packages 中的 protobuf（来自 onnx）（3.19.1）
已满足要求满意：c：\ users \ alessandro \ miniconda3 \ envs \ pt190 \ lib \ site-packages中的六个（来自onnx）（1.16.0）
import netron
import time
from IPython.display import IFrame
def show_netron(model_path, port):
time.sleep(3.)
netron.start(model_path, address=("localhost", port), browse=False)
return IFrame(src=f"http://localhost:{port}/", width="100%", height=400)
一般来说，标准 ONNX opset 不支持表示低于 8b 的量化。 此外，ONNX QOp 表示需要在层的一部分设置输出量化器。
在最近引入的 QDQ 表示形式（从 0.8 版本开始 Brevitas 开始支持）中，始终具有输出量化器的约束得到了放松，它仅使用 QuantizeLinear 和 DequantizeLinear 来表示量化，但即使有这种支持，仍然存在 仅限于 8b 量化。
导出到自定义量化ONNX（QONNX）
作为替代方案，我们可以将其导出到 QONNX，这是 Brevitas 定义的自定义 ONNX 方言，支持可以捕获这些信息的自定义量化运算符：
torch.manual_seed(0)
from brevitas.export import export_qonnx
from brevitas.quant import Int8WeightPerTensorFloat, Int8ActPerTensorFloat, Int16Bias
float_inp = torch.randn(1, 2, 5)
quant_conv_4b8b = QuantConv1d(
2, 4, 3, bias=True, weight_bit_width=4,
input_quant=Int8ActPerTensorFloat,
output_quant=Int8ActPerTensorFloat,
bias_quant=Int16Bias)
output_path = 'brevitas_onnx_conv4b8b.onnx'
export_qonnx(quant_conv_4b8b, input_t=float_inp, export_path=output_path)
通过这种方式，支持任意比例、零点和位宽，不再有上述8b的限制。
上面显示的自定义格式可以集成到基于 ONNX 的工具链中，例如 它由我们自己的 FINN 工具链支持，用于低精度数据流风格的定制 FPGA 实现，并且将成为与 TVM 直接集成的起点。
导出到 TorchScript 量化后端
二、QuantTensor 和 QuantConv2d 概述
QuantTensor是Brevitas中基本的数据结构，QuantConv2d是Brevitas中一个典型的层。QuantConv2d 是 QuantWeightBiasInputOutputLayer 的实例（通常作为 QuantWBIOL 导入，这意味着其支持权重、输入、输出和偏置的量化。QuantWBIOL的其他实例还有QuantLinear, QuantConv1d, QuantConvTranspose1d 和QuantConvTranspose2d，都遵循相同的原则。
观察QuantConv2d的初始化方法__init__：
import inspect
from brevitas.nn import QuantConv2d
from brevitas.nn import QuantIdentity
from IPython.display import Markdown, display
def pretty_print_source(source):
display(Markdown('```python\n' + source + '\n```'))
source = inspect.getsource(QuantConv2d.__init__)
pretty_print_source(source)
QuantConv2d是Conv2d和QuantWBIOL的共同实例，它的初始化方法公开了 Conv2d 的常用参数，以及：支持相同填充的额外标志padding_type；四个量化标志——weight_quant，bias_quant，input_quant和output_quant；输出QuantTensor的标志return_quant_tensor。
默认情况下weight_quant = Int8weightPerTensorFloat，其余bias_quant，input_quant和output_quant默认为None。这意味着默认情况下，权重被量化为带有每个张量浮点比例因子的 8 位有符号整数（ONNX 标准 opset 等采用的一种非常常见的量化类型），而偏差、输入和输出的量化是 禁用。 我们可以在运行时通过示例轻松验证所有这些：
default_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=False)
print(f'Is weight quant enabled: {default_quant_conv.is_weight_quant_enabled}')
print(f'Is bias quant enabled: {default_quant_conv.is_bias_quant_enabled}')
print(f'Is input quant enabled: {default_quant_conv.is_input_quant_enabled}')
print(f'Is output quant enabled: {default_quant_conv.is_output_quant_enabled}')
Is weight quant enabled: True
Is bias quant enabled: False
Is input quant enabled: False
Is output quant enabled: False
如果我们现在尝试传入一个随机浮点张量作为输入，正如预期的那样，我们会得到卷积的输出：
import torch
out = default_quant_conv(torch.randn(1, 2, 5, 5))
out
tensor([[[[-0.2594, 0.5392, 0.5916],
[ 0.3493, 0.6813, 0.2499],
[ 1.3732, 0.1229, -0.0084]],
[[ 0.0031, -0.1702, 0.1069],
[-0.8181, -0.8056, 0.0385],
[-0.4738, 0.0589, 0.1278]],
[[-0.1718, -0.1162, -0.1526],
[-0.9903, -0.3541, 0.1645],
[ 0.0557, -0.4458, -0.2080]]]], grad_fn=<ThnnConv2DBackward0>)
在这种情况下，我们正在计算未量化的输入张量和量化的权重之间的卷积，因此输出通常是未量化的。
当QuantConv2d禁用所有的量化时，它就相当于一个普通的Conv2d，可以验证如下：
from torch.nn import Conv2d
torch.manual_seed(0) # set a seed to make sure the random weight init is reproducible
disabled_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, weight_quant=None)
torch.manual_seed(0) # reproduce the same random weight init as above
float_conv = Conv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=False)
inp = torch.randn(1, 2, 5, 5)
assert torch.isclose(disabled_quant_conv(inp), float_conv(inp)).all().item()
Brevitas 允许用户尽可能自由地进行量化实验，这意味着量化和非量化值之间的计算被认为是合法的。 这允许用户将 Brevitas 层与 Pytorch 层混合使用，几乎没有任何限制。
为了实现这一点，量化值通常以反量化格式（存储时量化形式，运算时反量化）表示，转换公式为：quant_value = (integer_value - zero_point) * scale。
1.QuantTensor
可以调用quant_weight()查看量化权重：
default_quant_conv.quant_weight()
QuantTensor(value=tensor([[[[-0.0790, 0.0503, -0.0934],
[-0.1149, -0.1903, -0.1329],
[-0.1813, 0.0108, 0.0593]],
[[ 0.0970, -0.0215, -0.0144],
[ 0.2280, 0.1239, -0.0090],
[ 0.1957, -0.2011, -0.0108]]],
[[[-0.0018, -0.1957, 0.1993],
[-0.0359, 0.1778, -0.1400],
[ 0.0916, 0.1059, 0.2173]],
[[-0.1670, 0.1939, -0.2191],
[-0.0215, 0.1688, -0.1383],
[-0.0449, -0.1185, 0.1742]]],
[[[-0.0808, -0.1652, -0.0233],
[-0.0700, 0.0467, -0.0485],
[ 0.1059, 0.1418, 0.1077]],
[[-0.0593, 0.0108, 0.0036],
[-0.1508, 0.0808, 0.1616],
[ 0.0144, -0.0287, -0.1365]]]], grad_fn=<MulBackward0>), scale=tensor(0.0018, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
结果返回的张量，在Brevitas中称之为QuantTensor。QuantTensor 是一种表示仿射量化张量及其所有元数据的方法，它的值是反量化格式的量化张量的值（就是量化后的值，他在运算的时候需要进行反量化再运算），其余还包括进行反量化操作的一些参数scale,zero_point,bit_width,signed和training(是否训练模式).
有效的QuantTensor正确地用值填充其所有字段，并遵循仿射量化不变量，即（考虑舍入误差）可以在和字段定义的区间内表示的整数。无效的则不会。可以调用is_valid验证是否有效：
assert default_quant_conv.quant_weight().is_valid
但在某些情况下可能会生成无效的 QuantTensor，这一点需要注意。假设我们有两个 QuantTensor 作为相同量化激活的输出，并且我们想要将它们加在一起：
from brevitas.quant_tensor import QuantTensor
quant_act = QuantIdentity(return_quant_tensor=True)
out_tensor_0 = quant_act(torch.randn(1,2,5,5))
out_tensor_1 = quant_act(torch.randn(1,2,5,5))
assert out_tensor_0.is_valid
assert out_tensor_1.is_valid
print(out_tensor_0.scale)
print(out_tensor_1.scale)
张量（0.0173，grad_fn = <DivBackward0>）
张量（0.0307，grad_fn = <DivBackward0>）
两个 QuantTensor 都是有效的，但由于默认情况下量化激活处于训练模式，因此它们的比例因子将会不同。值得注意的是，评估时的行为是不同的，其中两个比例因子是必须相同的。
out_tensor = out_tensor_0 + out_tensor_1
out_tensor
QuantTensor(值=张量([[[[ 0.9489, -0.9111, -0.0536, 0.5788, 0.3645],
[ 0.3401, 1.4325, 0.6498, 0.6411, -1.4390],
[-1.9029, 0.7012, 0.1591, 1.9235, 0.5883],
[ -2.7258, 2.5330, 0.9165, -0.0820, 3.4148],
[-0.3651, 1.0164, 0.9567, -0.2758, -1.1376]],
[[-0.2414, 2.2111, -1.9124, -2.3814, -0.88 05]，
[1.3191，- 0.8965, -0.2048, -3.8113, 1.1142],
[-0.3381, -0.2238, 1.2661, 0.0068, 0.2567],
[ 0.0731, -0.4280, 0.0909, 0.0875, -1.6851],
[-0.77 44、-1.4127、-0.8143、1.3557 ，-0.2802]]]]，
grad_fn=<AddBackward0>），scale=张量（0.0240，grad_fn=<DivBackward0>），zero_point=张量（0.），bit_width=张量（9.），signed_t=张量（True） ,training_t=张量(真))
因为我们将它们都设置training为True，所以即使它们具有不同的比例因子，我们也可以对它们求和。输出 QuantTensor 将具有正确的bit_width和scale，该scale是两个原始比例因子的平均值。这仅在训练时完成，以便传播梯度信息，但结果是生成的QuantTensor不再有效：
assert not out_tensor.is_valid
QuantTensor 实现 __torch_function__ 来处理torch 函数运算符（例如 torch.nn.function 下的操作）的调用。有些操作（量化不相关）可以保持QuantTensor的有效性，如：max-pooling：
import torch
quant_identity = QuantIdentity(return_quant_tensor=True)
quant_tensor = quant_identity(torch.randn(1, 3, 4, 4))
torch.nn.functional.max_pool2d(quant_tensor, kernel_size=2, stride=2)
QuantTensor(value=tensor([[[[1.5800, 1.0157],
[1.4445, 0.8577]],
[[0.5643, 1.2414],
[1.0383, 0.9028]],
[[0.5191, 0.6546],
[2.1442, 0.5868]]]], grad_fn=<MaxPool2DWithIndicesBackward0>), scale=tensor(0.0226, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
有些运算是量化相关的，将会导致输出的QuantTensor衰减为普通的torch.Tensor：
torch.tanh(quant_tensor)
tensor([[[[-0.4943, -0.9938, -0.9073, 0.7681],
[-0.3262, 0.9186, 0.1786, 0.3659],
[ 0.7489, 0.8946, -0.0451, -0.5594],
[-0.1346, -0.4943, -0.4770, 0.6951]],
[[ 0.0676, 0.5111, 0.4943, 0.8459],
[-0.8990, -0.9426, 0.0676, -0.7945],
[-0.9220, 0.0676, -0.5594, 0.6321],
[-0.0676, 0.7772, 0.7177, -0.4414]],
[[ 0.4770, 0.2220, 0.0676, 0.5747],
[-0.0451, -0.6710, -0.4594, -0.3462],
[ 0.9729, -0.7177, -0.5896, -0.5276],
[-0.0900, 0.8852, 0.5276, -0.4414]]]], grad_fn=<TanhBackward0>)
2.输入量化
如果要QuantConv2d 输出是一个QuantTensor,则需要输入和权重都被量化。我们可以给input_quant设置一个量化器，下边的例子使用带有每个张量浮点比例因子的带符号 8 位量化器：
from brevitas.quant.scaled_int import Int8ActPerTensorFloat
input_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=False,
input_quant=Int8ActPerTensorFloat, return_quant_tensor=True)
out_tensor = input_quant_conv(torch.randn(1, 2, 5, 5))
out_tensor
QuantTensor(value=tensor([[[[ 0.9693, -0.9431, 0.2459],
[ 0.5416, 0.9037, -0.5278],
[-0.6207, -1.3578, -0.4815]],
[[ 0.4551, -1.4065, 0.8889],
[-0.3393, 0.0803, -0.1748],
[-0.0977, 0.6284, -0.7193]],
[[ 0.3655, 0.7626, -0.2634],
[-0.3453, 0.3349, 0.1923],
[ 0.5993, -0.9579, 0.3557]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[3.2208e-05]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))
内部发生的情况是，传递给 input_quant_conv 的输入张量在传递给卷积运算符之前被量化。这表明着我们现在正在计算两个量化张量之间的卷积，意味着操作的输出也是量化的。 正如预期的那样，out_tensor 被标记为有效。
另一件需要注意的重要事情是out_tensor 的 bit_width 字段的位宽为 21 位。在 Brevitas 中，给定输入和权重张量的大小及其位宽，21 是表示可以生成的最大可能输出值所需的位宽。 这确保了仿射量化不变量始终得到遵守。
我们可以通过直接传递 QuantTensor 作为输入来获得类似的结果。 在这个例子中，我们直接定义一个QuantTensor，但它也可以是前一层的输出。
from brevitas.quant_tensor import QuantTensor
scale = 0.0001
bit_width = 8
zero_point = 0.
int_value = torch.randint(low=- 2 ** (bit_width - 1), high=2 ** (bit_width - 1) - 1, size=(1, 2, 5, 5))
quant_value = (int_value - zero_point) * scale
quant_tensor_input = QuantTensor(
quant_value,
scale=torch.tensor(scale),
zero_point=torch.tensor(zero_point),
bit_width=torch.tensor(float(bit_width)),
signed=True,
training=True)
quant_tensor_input
QuantTensor(value=tensor([[[[ 5.7000e-03, 2.5000e-03, -1.2400e-02, -7.2000e-03, 3.7000e-03],
[-2.3000e-03, 7.0000e-04, -1.2700e-02, 5.2000e-03, 4.0000e-04],
[-7.9000e-03, 9.5000e-03, 6.6000e-03, 5.4000e-03, 2.5000e-03],
[ 1.1100e-02, 2.4000e-03, 1.0000e-02, -3.7000e-03, 7.2000e-03],
[-1.1500e-02, -5.8000e-03, -9.3000e-03, 1.0000e-02, 3.5000e-03]],
[[-6.8000e-03, 1.1500e-02, -1.0600e-02, -1.5000e-03, -1.9000e-03],
[ 2.9000e-03, 9.5000e-03, 7.2000e-03, -3.7000e-03, 7.7000e-03],
[-2.4000e-03, -8.9000e-03, -1.2000e-02, -8.1000e-03, 7.2000e-03],
[-1.1300e-02, -9.7000e-03, -1.0000e-03, 1.0100e-02, 3.8000e-03],
[-1.1900e-02, 6.9000e-03, 8.3000e-03, 1.0000e-04, -6.9000e-03]]]]), scale=tensor(1.0000e-04), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
将quant_tensor_input传入到return_quant_conv，实际上会得到有效的QuantTensor并且bit_width = 21：
return_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=False, return_quant_tensor=True)
out_tensor = return_quant_conv(quant_tensor_input)
QuantTensor(value=tensor([[[[ 0.0085, 0.0066, 0.0050],
[-0.0038, -0.0009, -0.0115],
[-0.0055, -0.0037, 0.0009]],
[[ 0.0015, -0.0027, -0.0079],
[-0.0034, -0.0060, 0.0043],
[-0.0008, 0.0052, -0.0033]],
[[-0.0015, 0.0082, -0.0038],
[-0.0021, 0.0004, -0.0054],
[-0.0021, -0.0079, 0.0013]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[1.8448e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))
同样可以向设置了input_quant参数的层传入QuantTensor，这样输入将会再次被量化：
input_quant_conv(quant_tensor_input)
QuantTensor(value=tensor([[[[-0.0035, -0.0037, -0.0050],
[ 0.0010, -0.0051, -0.0027],
[-0.0010, 0.0047, 0.0017]],
[[ 0.0021, 0.0002, 0.0027],
[ 0.0028, 0.0002, -0.0044],
[ 0.0008, -0.0052, -0.0024]],
[[ 0.0010, -0.0052, -0.0011],
[-0.0018, 0.0024, 0.0011],
[-0.0001, 0.0039, 0.0035]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[1.7410e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(21.), signed_t=tensor(True), training_t=tensor(True))
3.输出量化
当启用输出量化时：
from brevitas.quant.scaled_int import Int8ActPerTensorFloat
output_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=False,
output_quant=Int8ActPerTensorFloat, return_quant_tensor=True)
out_tensor = output_quant_conv(torch.randn(1, 2, 5, 5))
out_tensor
QuantTensor(value=tensor([[[[ 0.2111, 0.4060, 0.3654],
[-0.7876, 0.8119, -0.9825],
[-0.5115, 0.3979, -0.3248]],
[[ 0.3816, 0.0568, -0.0812],
[ 1.0312, -0.7876, 0.8038],
[-0.3491, -0.4141, 0.0650]],
[[-0.5846, -0.4222, -0.0731],
[-0.7389, 0.5034, -0.2517],
[-0.1624, -0.4385, 0.7308]]]], grad_fn=<MulBackward0>), scale=tensor(0.0081, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
输出是一个有效的QuantTensor，并且与仅设置输入量化有所不同。
- 只设置输出量化，则会输入浮点张量，因此内部计算量化张量和非量化张量的卷积，然后将结果进行量化；
- 只设置输出量化，我们可以看到结果中的
bit_width为8，符合量化器的设置。
4.偏置量化
偏差量化需要满足输入是一个QuantTensor，因为偏置的scale需要输入和权重的scale来计算。若仅仅设置QuantConv2d的bias_quant=Int8Bias，未进行任何的输入量化，则会报错：
from brevitas.quant.scaled_int import Int8Bias
bias_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,
bias_quant=Int8Bias, return_quant_tensor=True)
bias_quant_conv(torch.randn(1, 2, 5, 5))
可以通过传入QuantTensor解决上述报错：
bias_quant_conv(quant_tensor_input)
QuantTensor(value=tensor([[[[ 0.0005, 0.0043, -0.0004],
[ 0.0005, 0.0106, 0.0012],
[ 0.0021, 0.0007, -0.0050]],
[[-0.0067, -0.0035, -0.0059],
[-0.0050, -0.0015, -0.0039],
[ 0.0015, 0.0028, -0.0008]],
[[-0.0051, -0.0050, 0.0060],
[-0.0015, 0.0037, 0.0071],
[ 0.0067, 0.0035, -0.0071]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[1.8108e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))
或者设置输入量化，然后传入torch.Tensor或QuantTensor：
input_bias_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,
input_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias, return_quant_tensor=True)
input_bias_quant_conv(torch.randn(1, 2, 5, 5))
QuantTensor(value=tensor([[[[-0.3825, 0.1371, 0.9135],
[-0.2016, 0.7495, -0.4071],
[-0.0755, 0.5283, 0.2388]],
[[ 0.0788, -0.3802, -0.2234],
[ 0.8678, -0.5546, 0.4408],
[-0.6788, 0.4422, 0.3007]],
[[ 0.4412, -0.3205, 1.0033],
[-0.0083, -0.3295, -0.2076],
[ 0.4417, -0.1046, -0.3493]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[3.8610e-05]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))
input_bias_quant_conv(quant_tensor_input)
QuantTensor(value=tensor([[[[ 0.0036, 0.0024, -0.0033],
[ 0.0050, 0.0080, -0.0014],
[-0.0036, -0.0080, -0.0029]],
[[ 0.0083, -0.0093, 0.0048],
[ 0.0035, 0.0015, -0.0011],
[-0.0003, 0.0067, 0.0013]],
[[-0.0009, -0.0019, 0.0039],
[ 0.0010, 0.0056, -0.0037],
[ 0.0091, -0.0095, 0.0054]]]], grad_fn=<ThnnConv2DBackward0>), scale=tensor([[[[1.8384e-07]]]], grad_fn=<MulBackward0>), zero_point=tensor(0.), bit_width=tensor(22.), signed_t=tensor(True), training_t=tensor(True))
注意到输出的bit_width = 22，这是因为，在最坏的情况下，将 21 位整数（添加偏置之前累加器的大小）与 8 位整数（量化偏置的大小）相加会得到 22 位整数。
现在让我们尝试启用输出量化而不是输入量化。 这并不能解决偏置量化的问题，因为输出量化是在添加偏置之后执行的：
output_bias_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,
output_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias, return_quant_tensor=True)
output_bias_quant_conv(torch.randn(1, 2, 5, 5))
并非所有场景都需要偏差量化来取决于输入的比例因子。在这些情况下，偏差可以按照与权重量化相同的方式进行量化，并且具有自己的比例因子。在 Brevitas 中，反映这种其他情况的预定义量化器是Int8BiasPerTensorFloatInternalScaling。在这种情况下，不需要有效的量化输入：
from brevitas.quant.scaled_int import Int8BiasPerTensorFloatInternalScaling
bias_internal_scale_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,
bias_quant=Int8BiasPerTensorFloatInternalScaling, return_quant_tensor=False)
bias_internal_scale_quant_conv(torch.randn(1, 2, 5, 5))
张量([[[[ 0.2152, 0.8346, 0.0746],
[-0.0738, -0.5212, 0.1019],
[-0.6004, 0.1500, -0.1453]],
[[-1.1551, -1.3458, -0.1312],
[ 0.2502, - 0.5267, 0.2412],
[-0.3556, -0.3289, -0.2276]],
[[-0.4599, -0.6094, 0.4682],
[-0.5064, -0.6768, -0.6638],
[ 0.0066, -0.3581, 0.2359]] ]] , grad_fn=<ThnnConv2DBackward0>)
需要注意的是，偏置量化在一些场景下可能会导致输出的zero_point发生变化。
- 在输出之上添加一个未量化的偏差（如下例）；
- 添加了一个用其自己的比例因子（例如量化器
Int8BiasPerTensorFloatInternalScaling）量化的偏差；
为了确保输出QuantTensor有效，输出的zero_point变为非0：
unquant_bias_input_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,
input_quant=Int8ActPerTensorFloat, return_quant_tensor=True)
out_tensor = unquant_bias_input_quant_conv(torch.randn(1, 2, 5, 5))
out_tensor
QuantTensor(值=张量([[[[-0.6879,-0.6632,-0.2411],
[0.2064,-0.7371,0.3910],
[0.9533,0.2994,0.6546]],
[[-0.4684,-0.4495,-0.5021],
[ 0.5738, 0.4199, -0.3380],
[ 0.6218, -0.0408, -0.8483]],
[[-0.5625, 0.1837, -1.0575],
[-1.2816, -0.4993, -0.3409],
[ 0.4556, -1.4269, 0.5369] ]]]、grad_fn=<ThnnConv2DBackward0>）、scale=张量（[[[[3.0975e-05]]]]、grad_fn=<MulBackward0>）、zero_point=张量（[[[[ 1276.0774]]、
[[- 3152.4585]]、
[[ 7320.2324]]]]、grad_fn=<DivBackward0>)、bit_width=张量(21.)、signed_t=张量(真)、training_t=张量(真))
最后，一般我们没必要返回一个QuantTensor，因为它需要额外的开销，这也是为什么默认情况下return_quant_tensor设置为False:
bias_input_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), bias=True,
input_quant=Int8ActPerTensorFloat, bias_quant=Int8Bias)
bias_input_quant_conv(torch.randn(1, 2, 5, 5))
tensor([[[[ 0.8357, 0.0733, 0.9527],
[ 0.1803, 0.2154, 0.7598],
[ 1.1121, -0.8728, 1.0039]],
[[ 0.7917, 1.0063, 0.6516],
[-0.1852, -0.7263, 0.0956],
[-0.1876, 0.2747, -0.1617]],
[[ 0.8299, 0.9934, -0.3821],
[ 0.4865, 0.9309, -0.7924],
[-0.4201, 0.2343, 0.1532]]]], grad_fn=<ThnnConv2DBackward0>)
尽管输出的张量没有显示量化信息，但事实上上例的输出确实是量化了的。
三、量化激活概述
在QuantConv2d之后加入QuantIdentity层或者QuantConv2d设置了output_quant = Int8ActPerTensorFloat会输出相同的结果。下例设置相同的输入，验证上述两种情况：
import torch
from brevitas.nn import QuantConv2d, QuantIdentity
from brevitas.quant.scaled_int import Int8ActPerTensorFloat
torch.manual_seed(0)
output_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), output_quant=Int8ActPerTensorFloat)
torch.manual_seed(0)
default_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3))
output_identity_quant = QuantIdentity()
inp = torch.randn(1, 2, 5, 5)
out_tensor1 = output_quant_conv(inp)
out_tensor2 = output_identity_quant(default_quant_conv(inp))
assert out_tensor1.isclose(out_tensor2).all().item()
如果启用输入量化，也会发生类似情形：
torch.manual_seed(0)
input_output_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3),
input_quant=Int8ActPerTensorFloat, output_quant=Int8ActPerTensorFloat)
torch.manual_seed(0)
default_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3))
input_identity_quant = QuantIdentity()
output_identity_quant = QuantIdentity()
inp = torch.randn(1, 2, 5, 5)
out_tensor1 = input_output_quant_conv(inp)
out_tensor2 = output_identity_quant(default_quant_conv(input_identity_quant(inp)))
assert out_tensor1.isclose(out_tensor2).all().item()
从算法的角度来看，两种不同的实现正在做同样的事情。然而，正如在后面的教程中会变得更加清晰的那样，目前在某些情况下，在导出为标准 ONNX 等格式时，选择一种样式而不是另一种样式可能会产生不同的结果。
当禁用量化时，Quant_这些层就会与普通的浮点层没有区别。例如QuantIdentity禁用量化，意味着它就是一个普通的identity函数：
disabled_quant_identity = QuantIdentity(act_quant=None)
(inp == disabled_quant_identity(inp)).all().item()
量化激活层也可以返回QuantTensor：
return_quant_identity = QuantIdentity(return_quant_tensor=True)
out_tensor = return_quant_identity(inp)
out_tensor
QuantTensor(value=tensor([[[[-0.4566, -0.5707, -0.5517, 0.5897, 1.5409],
[ 0.5136, -0.5897, -0.5707, 0.1902, -0.0761],
[-0.4946, -1.5029, -0.1902, 0.4376, 1.3317],
[-1.6361, 2.0736, 1.7122, 2.3780, -1.1224],
[-0.3234, -1.0844, -0.0761, -0.0951, -0.7610]],
[[-1.5980, 0.0190, -0.7419, 0.1902, 0.6278],
[ 0.6468, -0.2473, -0.5327, 1.1605, 0.4376],
[-0.7990, -1.2936, -0.7419, -1.3127, -0.2283],
[-2.4351, -0.0761, 0.2283, 0.7990, -0.1902],
[-0.3615, -1.2175, -0.6278, -0.4566, 1.9214]]]],
grad_fn=<MulBackward0>), scale=tensor(0.0190, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
即使输入QuantTensor，禁用量化的QuantIdentity也和identity函数的作用一致。但是，无论是否设置return_quant_tensor，量化数据可能会被删除，并且输入QuantTensor可能会返回torch.Tensor(我观察结果没啥区别啊，不知道他说的什么意思):
out_torch_tensor = disabled_quant_identity(out_tensor)
out_torch_tensor
tensor([[[[-0.4566, -0.5707, -0.5517, 0.5897, 1.5409],
[ 0.5136, -0.5897, -0.5707, 0.1902, -0.0761],
[-0.4946, -1.5029, -0.1902, 0.4376, 1.3317],
[-1.6361, 2.0736, 1.7122, 2.3780, -1.1224],
[-0.3234, -1.0844, -0.0761, -0.0951, -0.7610]],
[[-1.5980, 0.0190, -0.7419, 0.1902, 0.6278],
[ 0.6468, -0.2473, -0.5327, 1.1605, 0.4376],
[-0.7990, -1.2936, -0.7419, -1.3127, -0.2283],
[-2.4351, -0.0761, 0.2283, 0.7990, -0.1902],
[-0.3615, -1.2175, -0.6278, -0.4566, 1.9214]]]],
grad_fn=<MulBackward0>)
return_disabled_quant_identity = QuantIdentity(act_quant=None, return_quant_tensor=True)
identity_out_tensor = return_disabled_quant_identity(out_tensor)
identity_out_tensor
QuantTensor(value=tensor([[[[-0.4566, -0.5707, -0.5517, 0.5897, 1.5409],
[ 0.5136, -0.5897, -0.5707, 0.1902, -0.0761],
[-0.4946, -1.5029, -0.1902, 0.4376, 1.3317],
[-1.6361, 2.0736, 1.7122, 2.3780, -1.1224],
[-0.3234, -1.0844, -0.0761, -0.0951, -0.7610]],
[[-1.5980, 0.0190, -0.7419, 0.1902, 0.6278],
[ 0.6468, -0.2473, -0.5327, 1.1605, 0.4376],
[-0.7990, -1.2936, -0.7419, -1.3127, -0.2283],
[-2.4351, -0.0761, 0.2283, 0.7990, -0.1902],
[-0.3615, -1.2175, -0.6278, -0.4566, 1.9214]]]],
grad_fn=<MulBackward0>), scale=tensor(0.0190, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(True), training_t=tensor(True))
以上各种设置同样适用于QuantReLU，不同的是QuantReLU先进行ReLU函数计算再进行量化，而QuantIdentity只是进行量化操作。此外，默认情况下QuantReLU采用Uint8ActPerTensorFloat量化器，意味着其量化输出是无符号的：
from brevitas.nn import QuantReLU
return_quant_relu = QuantReLU(return_quant_tensor=True)
return_quant_relu(inp)
QuantTensor(value=tensor([[[[0.0000, 0.0000, 0.0000, 0.5974, 1.5402],
[0.5041, 0.0000, 0.0000, 0.1867, 0.0000],
[0.0000, 0.0000, 0.0000, 0.4481, 1.3255],
[0.0000, 2.0817, 1.7083, 2.3804, 0.0000],
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],
[[0.0000, 0.0187, 0.0000, 0.1867, 0.6254],
[0.6348, 0.0000, 0.0000, 1.1668, 0.4387],
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.2334, 0.7935, 0.0000],
[0.0000, 0.0000, 0.0000, 0.0000, 1.9230]]]], grad_fn=<MulBackward0>), scale=tensor(0.0093, grad_fn=<DivBackward0>), zero_point=tensor(0.), bit_width=tensor(8.), signed_t=tensor(False), training_t=tensor(True))
QuantReLU，和QuantIdentity一样，不同于其他的非线性量化层，即使QuantReLU量化被禁用，它也保留输入 QuantTensor的数据。
return_disabled_quant_relu = QuantReLU(act_quant=None, return_quant_tensor=True)
relu_out_tensor = return_disabled_quant_relu(out_tensor)
assert relu_out_tensor.is_valid==True
assert relu_out_tensor.scale == out_tensor.scale
assert relu_out_tensor.zero_point == out_tensor.zero_point
assert relu_out_tensor.bit_width == out_tensor.bit_width
这不适用于有些层，比如QuantSigmoid：
from brevitas.nn import QuantSigmoid
return_disabled_quant_sigmoid = QuantSigmoid(act_quant=None, return_quant_tensor=True)
sigmoid_out_tensor = return_disabled_quant_sigmoid(out_tensor)
sigmoid_out_tensor
QuantTensor(value=(tensor([[[[0.3878, 0.3611, 0.3655, 0.6433, 0.8236],
[0.6257, 0.3567, 0.3611, 0.5474, 0.4810],
[0.3788, 0.1820, 0.4526, 0.6077, 0.7911],
[0.1630, 0.8883, 0.8471, 0.9151, 0.2456],
[0.4198, 0.2527, 0.4810, 0.4762, 0.3184]],
[[0.1683, 0.5048, 0.3226, 0.5474, 0.6520],
[0.6563, 0.4385, 0.3699, 0.7614, 0.6077],
[0.3102, 0.2152, 0.3226, 0.2120, 0.4432],
[0.0805, 0.4810, 0.5568, 0.6898, 0.4526],
[0.4106, 0.2284, 0.3480, 0.3878, 0.8723]]]],
grad_fn=<SigmoidBackward0>), None, None, None), scale=None, zero_point=None, bit_width=None, signed_t=None, training_t=tensor(True))
需要始终记住的一点是，量化激活层的非线性总是在输入的反量化表示上调用。例如，首先使用一个无符号移位量化器ShiftedUint8ActPerTensorFloat对一个浮点torch.Tensor进行量化，具有零点，使得其输出的整数表示是非负的。然后，将这个张量输入到禁用量化的QuantReLU。QuantReLU 的整数形式的输入是无符号的这一事实并不意味着不会对QuantReLU产生影响，因为 ReLU 是在反量化表示上调用的，其中包括正值和负值：
from brevitas.quant.shifted_scaled_int import ShiftedUint8ActPerTensorFloat
shifted_quant_identity = QuantIdentity(act_quant=ShiftedUint8ActPerTensorFloat, return_quant_tensor=True)
return_disabled_quant_relu = QuantReLU(act_quant=None, return_quant_tensor=True)
return_disabled_quant_relu(shifted_quant_identity(inp))
QuantTensor(value=tensor([[[[0.0000, 0.0000, 0.0000, 0.5854, 1.5485],
[0.5099, 0.0000, 0.0000, 0.1888, 0.0000],
[0.0000, 0.0000, 0.0000, 0.4532, 1.3219],
[0.0000, 2.0772, 1.6996, 2.3794, 0.0000],
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000]],
[[0.0000, 0.0189, 0.0000, 0.1888, 0.6232],
[0.6421, 0.0000, 0.0000, 1.1708, 0.4343],
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.2266, 0.7931, 0.0000],
[0.0000, 0.0000, 0.0000, 0.0000, 1.9262]]]], grad_fn=<ReluBackward0>), scale=tensor(0.0189, grad_fn=<DivBackward0>), zero_point=tensor(129., grad_fn=<SWhereBackward0>), bit_width=tensor(8.), signed_t=tensor(False), training_t=tensor(True))
接下来考虑一种非常常见的场景——一个QuantConv2d后边跟着一个ReLU或QuantReLU。尤其是设置了输出量化的QuantConv2d后接一个ReLU：
torch.manual_seed(0)
output_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3), output_quant=Int8ActPerTensorFloat)
torch.relu(output_quant_conv(inp))
tensor([[[[0.0000, 0.0000, 0.0000],
[1.3134, 1.2557, 1.0392],
[0.4186, 0.0000, 0.0000]],
[[0.7361, 0.5340, 0.8516],
[0.2887, 0.3175, 0.0000],
[0.8949, 1.6743, 0.0722]],
[[0.0000, 0.0000, 0.0289],
[0.0000, 0.0000, 0.2021],
[0.0000, 0.0000, 0.4907]]]], grad_fn=<ReluBackward0>)
将上述情况与默认设置的QuantConv2d（输出量化禁用）后接QuantReLU（输出量化启用）这种情况进行对比：
torch.manual_seed(0)
default_quant_conv = QuantConv2d(
in_channels=2, out_channels=3, kernel_size=(3,3))
default_quant_relu = QuantReLU()
default_quant_relu(default_quant_conv(inp))
tensor([[[[0.0000, 0.0000, 0.0000],
[1.3078, 1.2555, 1.0397],
[0.4185, 0.0000, 0.0000]],
[[0.7454, 0.5427, 0.8566],
[0.2943, 0.3269, 0.0000],
[0.8893, 1.6674, 0.0785]],
[[0.0065, 0.0000, 0.0262],
[0.0000, 0.0000, 0.1962],
[0.0000, 0.0000, 0.4839]]]], grad_fn=<MulBackward0>)
输出结果接近但不完全相同。
在第一种情况下，我们使用 8 位有符号量化器对 QuantConv2d 的输出进行量化，然后将其传递给 ReLU，这意味着有符号量化器覆盖的数值范围的一半现在丢失了，并且通过所有实际方法， 输出现在可以被视为 7 位无符号数（尽管没有明确标记为这样）。在第二种情况下，我们在 ReLU 之后执行无符号 8 位量化。 由于量化器覆盖的范围现在仅包含非负数，因此我们不会像前一种情况那样浪费一点。(个人理解：第一种情况是先量化再ReLU，很明显量化包含了 QuantConv2d 所有的输出；第二种情况是先ReLU再量化，量化的数据只包含了 QuantConv2d输出的大于0的部分，因而结果接近但不完全相同)
关于一些预制的激活量化器，例如Uint8ActPerTensorFloat，ShiftedUint8ActPerTensorFloat和Int8ActPerTensorFloat，预示着下一个教程的一些主题。为了最大限度地减少用户交互，Brevitas 通过收集多个训练步骤（默认为 30 个）的统计数据来初始化比例和零点。这可以被视为一种非常基本的校准步骤，尽管它通常发生在训练期间并且已经启用量化。这些统计数据以指数移动平均值的形式累积，在收集阶段结束时用于初始化学习参数。在收集阶段，量化器在 train() 和 eval() 模式之间的行为有所不同。 在 train() 模式下，返回该特定批次的统计信息。 在 eval() 模式下，返回指数移动平均线。 收集阶段结束后，在两种执行模式下都会返回学习到的参数。 我们可以通过一个例子很容易地观察到这种行为。 我们首先定义一个量化激活和两个随机输入张量：
quant_identity = QuantIdentity(return_quant_tensor=True)
inp1 = torch.randn(3, 3)
inp2 = torch.randn(3, 3)
比较 train() 和 eval()模式之间两个张量的输出比例因子，一般来说， train()模式下的情况是不同的， eval()模式下的都是一样的。
out1_train = quant_identity(inp1)
out2_train = quant_identity(inp2)
assert not out1_train.scale.isclose(out2_train.scale).item()
False
quant_identity.eval()
out1_eval = quant_identity(inp1)
out2_eval = quant_identity(inp2)
assert out1_eval.scale.isclose(out2_eval.scale).item()
True
默认情况下，唯一例外的是QuantHardTanh层。这是因为 torch.nn.HardTanh 的接口已经要求用户手动指定 min_val 和 max_val。因此 Brevitas 在启用或禁用量化时都会保留这一点。 启用量化后，默认情况下这些值用于初始化，但随后会学习范围。 让我们看一个例子：
from brevitas.nn import QuantHardTanh
QuantHardTanh()
---------------------------------------------------------------------------
DependencyError Traceback (most recent call last)
<ipython-input-18-8145d2f87fcb> in <module>
1 from brevitas.nn import QuantHardTanh
2
----> 3 QuantHardTanh()
c:\brevitas_fx\src\brevitas\nn\quant_activation.py in __init__(self, act_quant, input_quant, return_quant_tensor, **kwargs)
117 act_quant=act_quant,
118 return_quant_tensor=return_quant_tensor,
--> 119 **kwargs)
120
121
c:\brevitas_fx\src\brevitas\nn\quant_layer.py in __init__(self, act_impl, passthrough_act, input_quant, act_quant, return_quant_tensor, **kwargs)
77 passthrough_act,
78 act_quant,
---> 79 **kwargs)
80
81 @property
c:\brevitas_fx\src\brevitas\nn\mixin\act.py in __init__(self, act_impl, passthrough_act, act_quant, **kwargs)
157 proxy_prefix='act_',
158 kwargs_prefix='',
--> 159 **kwargs)
160
161 @property
c:\brevitas_fx\src\brevitas\nn\mixin\base.py in __init__(self, quant, proxy_protocol, none_quant_injector, proxy_prefix, kwargs_prefix, **kwargs)
98 quant_injector = quant
99 quant_injector = quant_injector.let(**filter_kwargs(kwargs_prefix, kwargs))
--> 100 quant = quant_injector.proxy_class(self, quant_injector)
101 else:
102 if not isinstance(quant, proxy_protocol):
c:\brevitas_fx\src\brevitas\proxy\runtime_quant.py in __init__(self, quant_layer, quant_injector)
108
109 def __init__(self, quant_layer, quant_injector):
--> 110 super(ActQuantProxyFromInjector, self).__init__(quant_layer, quant_injector)
111 self.is_passthrough_act = _is_passthrough_act(quant_injector)
112
c:\brevitas_fx\src\brevitas\proxy\quant_proxy.py in __init__(self, quant_layer, quant_injector, export_mode, export_handler)
74 # Use a normal list and not a ModuleList since this is a pointer to parent modules
75 self.tracked_module_list = []
---> 76 self.add_tracked_module(quant_layer)
77 self.export_handler = export_handler
78 self.export_mode = export_mode
c:\brevitas_fx\src\brevitas\proxy\quant_proxy.py in add_tracked_module(self, module)
130 self.tracked_module_list.append(module)
131 self.update_tracked_modules()
--> 132 self.init_tensor_quant()
133 else:
134 raise RuntimeError("Trying to add None as a parent module.")
c:\brevitas_fx\src\brevitas\proxy\runtime_quant.py in init_tensor_quant(self)
120
121 def init_tensor_quant(self):
--> 122 tensor_quant = self.quant_injector.tensor_quant
123 act_impl = self.quant_injector.act_impl
124 is_act_enabled = _is_act_enabled(act_impl, tensor_quant)
[... skipping hidden 1 frame]
DependencyError: 'Int8ActPerTensorFloatMinMaxInit' can not resolve attribute 'max_val' while building 'scaling_init_impl'
由于没有设置 min_val 和 max_val，出现了以上报错。修改后：
quant_hard_tanh = QuantHardTanh(max_val=1.0, min_val=-1.0, return_quant_tensor=True)
该层现已正确初始化。 我们可以看到 train() 和 eval() 模式之间的输出比例因子都是相同的：
out1_train = quant_hard_tanh(inp1)
quant_hard_tanh.eval()
out2_eval = quant_hard_tanh(inp2)
assert out1_train.scale.isclose(out2_eval.scale).item()
True
最后，在 Brevitas 中，混合使用是完全合法的，并且提倡这么做。例如，设置了act_quant=Int8ActPerTensorFloatMinMaxInit的QuantIdentity层相当于默认的QuantHardTanh，或者相反启用了act_quant=Int8ActPerTensorFloatMinMaxInit的QuantHardTanh相当于默认的QuantIdentity。这是因为当设置不同的量化器时，同一层可以接受不同的关键字参数。 因此，带有 act_quant=Int8ActPerTensorFloatMinMaxInit 的 QuantIdentity 将需要参数 min_val 和 max_val，就像默认的 QuantHardTanh 一样。
四、量化器解析
Brevitas中的量化器是brevitas.inject.ExtendedInjector的子类，带有tensor_quant属性，该属性指向实现量化的torch模块的实例。量化器从 brevitas.quant 导入并传递到量化层：
from brevitas.inject import ExtendedInjector
from brevitas.quant.scaled_int import Int8ActPerTensorFloat
issubclass(Int8ActPerTensorFloat, ExtendedInjector)
True
Int8ActPerTensorFloat.tensor_quant
RescalingIntQuant(
(int_quant): IntQuant(
(float_to_int_impl): RoundSte()
(tensor_clamp_impl): TensorClamp()
(delay_wrapper): DelayWrapper(
(delay_impl): _NoDelay()
)
)
(scaling_impl): ParameterFromRuntimeStatsScaling(
(stats_input_view_shape_impl): OverTensorView()
(stats): _Stats(
(stats_impl): AbsPercentile()
)
(restrict_scaling): _RestrictValue(
(restrict_value_impl): FloatRestrictValue()
)
(clamp_scaling): _ClampValue(
(clamp_min_ste): ScalarClampMinSte()
)
(restrict_inplace_preprocess): Identity()
(restrict_preprocess): Identity()
)
(int_scaling_impl): IntScaling()
(zero_point_impl): ZeroZeroPoint(
(zero_point): StatelessBuffer()
)
(msb_clamp_bit_width_impl): BitWidthConst(
(bit_width): StatelessBuffer()
)
)
注意，量化器是子类而不是实例，要搞明白为什么，需要先理解ExtendedInjector是什么，为什么要用它。
1. 使用自动装配依赖注入进行量化
Pytorch 因其简单的类似 numpy 的“按运行定义”执行模型而大受欢迎。然而，当涉及到应用量化时，这种编程风格会带来问题。
许多量化方法依赖于基于state_dict原始浮点模型（用 Pytorch 术语）做出决策，以通过量化进行微调。然而，当我们在 Pytorch中实例化一个模型时，我们无法当场知道几行代码后是否会加载 state_dict。然而，因为 Pytorch 是按运行定义的，所以我们需要我们的模型在 state_dict 可能加载之前和之后一致地工作。 在传统场景中，这不会造成问题。 然而，在循环中进行量化时，量化器的定义方式可能会在加载预训练的 state_dict 之前和之后发生变化。
这意味着我们需要一种方法来定义我们的量化模型，以便它可以在 state_dict 发生变化时做出适当的反应。 在仅使用 Python的世界中，这并不会太难。 然而，为了减轻量化感知训练对性能的影响，Brevitas 扩展使用 Pytorch 的 JIT 编译器来实现 Python 的自定义子集 TorchScript。 这意味着在大多数情况下，当加载 state_dict 时，我们需要重新编译模型的部分内容。 由于编译通常是一个有损过程，因此 TorchScript 组件不能简单地根据新的输入信息重新编译自身。
然后，我们需要一种方法来声明量化方法，以便在 state_dict 发生变化时可以重新初始化并进行 JIT 编译。 因为我们想要支持任意复杂的用户定义的量化算法，所以该方法必须是通用的，即它不能依赖于所实现的量化算法的细节。
使用 ExtendedInjector实现量化器是一种方法。 具体来说，ExtendInjector 从旧版本 (0.2.1) 的优秀依赖注入库依赖项中扩展了 Injector，并支持一些特定于 Brevitas 需求的额外功能。
Injector（和 ExtendedInjector）允许获取可能非常复杂的交织对象图，并将其转换为能够通过将变量名称与参数名称匹配来自动组装的变量的平面列表。 这种技术通常被称为自动装配依赖注入。
在 Brevitas 的背景下，目标是收集有助于量化实现的所有模块和超参数，以便它们可以根据需要自动重新组装。 这个过程产生的是一个tensor_quant对象。
2. 一个实际的例子：二进制量化
通常实现量化的组件都可以在brevitas.core中找到。之前提到，Brevitas 大量使用 TorchScript。特别是，在 brevitas.core 下找到的所有组件都被实现为可以组装在一起的 ScriptModule。实现二值化的核心ScriptModule可以在brevitas.core.quant下找到：
import inspect
from IPython.display import Markdown, display
def pretty_print_source(source):
display(Markdown('```python\n' + source + '\n```'))
from brevitas.core.quant import BinaryQuant
source = inspect.getsource(BinaryQuant)
pretty_print_source(source)
class BinaryQuant(brevitas.jit.ScriptModule):
"""
ScriptModule that implements scaled uniform binary quantization of an input tensor.
Quantization is performed with :func:`~brevitas.function.ops_ste.binary_sign_ste`.
Args:
scaling_impl (Module): Module that returns a scale factor.
quant_delay_steps (int): Number of training steps to delay quantization for. Default: 0
Returns:
Tuple[Tensor, Tensor, Tensor, Tensor]: Quantized output in de-quantized format, scale, zero-point, bit_width.
Examples:
>>> from brevitas.core.scaling import ConstScaling
>>> binary_quant = BinaryQuant(ConstScaling(0.1))
>>> inp = torch.Tensor([0.04, -0.6, 3.3])
>>> out, scale, zero_point, bit_width = binary_quant(inp)
>>> out
tensor([ 0.1000, -0.1000, 0.1000])
>>> scale
tensor(0.1000)
>>> zero_point
tensor(0.)
>>> bit_width
tensor(1.)
Note:
Maps to quant_type == QuantType.BINARY == 'BINARY' == 'binary' when applied to weights in higher-level APIs.
Note:
Set env variable BREVITAS_JIT=1 to enable TorchScript compilation of this module.
"""
def __init__(self, scaling_impl: Module, quant_delay_steps: int = 0):
super(BinaryQuant, self).__init__()
self.scaling_impl = scaling_impl
self.bit_width = BitWidthConst(1)
self.zero_point = StatelessBuffer(torch.tensor(0.0))
self.delay_wrapper = DelayWrapper(quant_delay_steps)
@brevitas.jit.script_method
def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
scale = self.scaling_impl(x)
y = binary_sign_ste(x) * scale
y = self.delay_wrapper(x, y)
return y, scale, self.zero_point(), self.bit_width()
实现非常简单。 除了允许将量化延迟一定数量的训练步骤（默认 = 0）的 quant_delay_steps 之外，BinaryQuant 接受的唯一其他参数是计算比例因子的实现。 bit_width 固定为 1，zero_point固定为 0。
我们选择名为 ParameterScaling 的 ScriptModule 作为比例因子实现，它通过用户定义的初始化实现学习参数。 它可以在 brevitas.core.scaling 下找到：
from brevitas.core.scaling import ParameterScaling
手动二进制量化
第一步，我们简单地使用 ParameterScaling 实例化 BinaryQuant，使用的scaling_init 等于 0.1，并在随机浮点输入张量上调用它：
import torch
manual_tensor_quant = BinaryQuant(scaling_impl=ParameterScaling(scaling_init=0.1))
manual_tensor_quant(torch.randn(4, 4))
(tensor([[ 0.1000, 0.1000, 0.1000, 0.1000],
[-0.1000, -0.1000, 0.1000, 0.1000],
[ 0.1000, -0.1000, 0.1000, -0.1000],
[ 0.1000, -0.1000, 0.1000, -0.1000]], grad_fn=<MulBackward0>),
tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>),
tensor(0.),
tensor(1.))
正如预期的那样，输入张量使用我们定义的比例因子进行二值化，但注意manual_tensor_quant返回的是tuple而不是QuantTensor。这是因为 TorchScript 对自定义数据结构的支持仍然相当有限，因此 QuantTensor 仅在 Python 世界的抽象中分配(意思就是不支持)。
使用 ExtendedInjector 进行二进制量化
通过ExtendedInjector声明一个tensor_quant：
from brevitas.inject import ExtendedInjector
class MyBinaryQuantizer(ExtendedInjector):
tensor_quant = BinaryQuant
scaling_impl=ParameterScaling
scaling_init=0.1
inj_tensor_quant = MyBinaryQuantizer.tensor_quant
inj_tensor_quant(torch.randn(4, 4))
(tensor([[-0.1000, 0.1000, -0.1000, 0.1000],
[ 0.1000, 0.1000, -0.1000, -0.1000],
[-0.1000, 0.1000, -0.1000, 0.1000],
[-0.1000, 0.1000, 0.1000, 0.1000]], grad_fn=<MulBackward0>),
tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>),
tensor(0.),
tensor(1.))
每当调用 MyBinaryQuantizer.tensor_quant 时，都会创建 BinaryQuant 的新实例。注意 MyBinaryQuantizer 的属性是如何设计为与彼此的参数名称匹配的（tensor_quant 除外，这是我们有兴趣从外部检索的内容）。
量化器的继承和组合
通过 Python 类表达量化器的优点还意味着我们可以利用继承和组合。例如，我们可以继承MyBinaryQuantizer并覆盖scaling_init新值：
class MyChildBinaryQuantizer(MyBinaryQuantizer):
scaling_init=1.0
child_inj_tensor_quant = MyChildBinaryQuantizer.tensor_quant
child_inj_tensor_quant(torch.randn(4, 4))
(tensor([[ 1., -1., 1., 1.],
[ 1., 1., -1., 1.],
[ 1., 1., 1., -1.],
[-1., 1., -1., -1.]], grad_fn=<MulBackward0>),
tensor(1., grad_fn=<AbsBinarySignGradFnBackward>),
tensor(0.),
tensor(1.))
或者我们可以通过将包含量化器不同部分的各种类组装在一起来利用组合：
class MyBinaryImpl(ExtendedInjector):
tensor_quant = BinaryQuant
class MyScalingImpl(ExtendedInjector):
scaling_impl=ParameterScaling
scaling_init=0.1
class MyComposedBinaryQuantizer(MyBinaryImpl, MyScalingImpl):
pass
comp_inj_tensor_quant = MyComposedBinaryQuantizer.tensor_quant
comp_inj_tensor_quant(torch.randn(4, 4))
(tensor([[ 0.1000, -0.1000, -0.1000, -0.1000],
[-0.1000, 0.1000, -0.1000, 0.1000],
[ 0.1000, -0.1000, 0.1000, 0.1000],
[-0.1000, 0.1000, -0.1000, 0.1000]], grad_fn=<MulBackward0>),
tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>),
tensor(0.),
tensor(1.))
将量化器和量化层连接
在我们将量化器传递到量化层（例如 QuantConv2d）之前，我们需要定义最后一个组件，即代理(proxy)。代理（位于 brevitas.proxy 下）是一个 nn.Module，充当量化器和量化层之间的接口。
虽然量化器主要存在于JIT域，但代理主要存在于 Python 域，因此可以提供更大的灵活性。每当加载新的 state_dict 时，代理都会返回 QuantTensor 并重新初始化量化器的输出。
代理特定于被量化的张量类型，如权重、偏差和激活。 为了方便起见，它们在属性 proxy_class下声明为量化器本身的一部分。例如，对于权重可以使用WeightQuantProxyFromInjector：
from brevitas.proxy import WeightQuantProxyFromInjector
class MyBinaryWeightQuantizer(MyBinaryQuantizer):
proxy_class = WeightQuantProxyFromInjector
现在可以使用MyBinaryWeightQuantizer作为图层的权重量化器：
from brevitas.nn import QuantConv2d
binary_weight_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MyBinaryWeightQuantizer)
quant_weight = binary_weight_quant_conv.quant_weight()
quant_weight
QuantTensor(value=tensor([[[[ 0.1000, 0.1000, -0.1000],
[-0.1000, 0.1000, -0.1000],
[ 0.1000, -0.1000, -0.1000]],
[[-0.1000, 0.1000, -0.1000],
[ 0.1000, -0.1000, 0.1000],
[-0.1000, -0.1000, 0.1000]],
[[ 0.1000, -0.1000, -0.1000],
[ 0.1000, 0.1000, -0.1000],
[-0.1000, -0.1000, 0.1000]]],
[[[ 0.1000, -0.1000, 0.1000],
[ 0.1000, -0.1000, -0.1000],
[ 0.1000, -0.1000, 0.1000]],
[[-0.1000, 0.1000, -0.1000],
[ 0.1000, 0.1000, 0.1000],
[-0.1000, -0.1000, -0.1000]],
[[ 0.1000, 0.1000, -0.1000],
[-0.1000, 0.1000, -0.1000],
[ 0.1000, 0.1000, 0.1000]]]], grad_fn=<MulBackward0>), scale=tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=None, training_t=tensor(True))
注意 QuantTensor 的格式不正确，因为signed_t=None。这意味着 quant_weight 不被认为是有效的，因为无法计算仿射量化不变量：
signed是在二进制量化的情况下必须由用户显式定义的属性之一。 这个想法是，它通知代理我们的量化器生成的值是否应该被视为有符号。 我们可以通过简单地在量化器中设置它来做到这一点：
class MySignedBinaryWeightQuantizer(MyBinaryWeightQuantizer):
signed = True
binary_weight_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MySignedBinaryWeightQuantizer)
signed_quant_weight = binary_weight_quant_conv.quant_weight()
signed_quant_weight
QuantTensor(value=tensor([[[[ 0.1000, 0.1000, -0.1000],
[-0.1000, -0.1000, 0.1000],
[ 0.1000, 0.1000, -0.1000]],
[[ 0.1000, 0.1000, 0.1000],
[ 0.1000, -0.1000, 0.1000],
[ 0.1000, -0.1000, -0.1000]],
[[-0.1000, 0.1000, 0.1000],
[ 0.1000, -0.1000, -0.1000],
[-0.1000, -0.1000, -0.1000]]],
[[[ 0.1000, 0.1000, 0.1000],
[ 0.1000, 0.1000, -0.1000],
[-0.1000, -0.1000, 0.1000]],
[[-0.1000, -0.1000, 0.1000],
[-0.1000, 0.1000, 0.1000],
[-0.1000, -0.1000, -0.1000]],
[[-0.1000, 0.1000, -0.1000],
[-0.1000, 0.1000, -0.1000],
[-0.1000, 0.1000, -0.1000]]]], grad_fn=<MulBackward0>), scale=tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
现在，quant_weight是有效的。
当我们想要添加或覆盖传递到图层的量化器的单个属性时，定义一个全新的量化器可能会过于冗长。 有一种更简单的语法可以实现相同的目标。 假设我们想要将signed 属性添加到MyBinaryQuantizer，就像我们刚才所做的那样。 我们也可以简单地执行以下操作：
small_scale_quant_conv = QuantConv2d(3, 2, (3,3), weight_quant=MyBinaryWeightQuantizer, weight_signed=True)
small_scale_quant_conv.quant_weight()
QuantTensor(value=tensor([[[[-0.1000, -0.1000, 0.1000],
[-0.1000, -0.1000, -0.1000],
[ 0.1000, -0.1000, 0.1000]],
[[-0.1000, 0.1000, -0.1000],
[-0.1000, 0.1000, 0.1000],
[ 0.1000, -0.1000, -0.1000]],
[[-0.1000, 0.1000, -0.1000],
[-0.1000, -0.1000, 0.1000],
[-0.1000, -0.1000, -0.1000]]],
[[[-0.1000, -0.1000, -0.1000],
[-0.1000, -0.1000, -0.1000],
[ 0.1000, 0.1000, -0.1000]],
[[-0.1000, -0.1000, 0.1000],
[-0.1000, 0.1000, -0.1000],
[ 0.1000, -0.1000, 0.1000]],
[[ 0.1000, 0.1000, -0.1000],
[ 0.1000, 0.1000, 0.1000],
[ 0.1000, -0.1000, 0.1000]]]], grad_fn=<MulBackward0>), scale=tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
将想要修改的参数名加上weight_前缀，并将其作为关键字参数传递给QuantConv2d。以weight_为前缀的关键字参数被设置为weight_quant的属性，可能会覆盖任何预先存在的值。相同的原则同样适用于input_, output_ and bias_。
将自定义量化器传递给 QuantIdentity
量化激活也类似如此：
from brevitas.proxy import ActQuantProxyFromInjector
from brevitas.nn import QuantIdentity
class MySignedBinaryActQuantizer(MyBinaryQuantizer):
proxy_class = ActQuantProxyFromInjector
signed = True
binary_relu = QuantIdentity(act_quant=MySignedBinaryActQuantizer, return_quant_tensor=True)
binary_relu(torch.randn(4, 4))
QuantTensor(value=tensor([[-0.1000, 0.1000, -0.1000, 0.1000],
[ 0.1000, 0.1000, 0.1000, 0.1000],
[-0.1000, 0.1000, 0.1000, 0.1000],
[-0.1000, -0.1000, 0.1000, -0.1000]], grad_fn=<MulBackward0>), scale=tensor(0.1000, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
因此，权重量化器和激活量化器之间并没有太大区别，它们只是由不同的代理包装。 此外，对于激活，传递关键字参数时不需要前缀。 例如，传入新值覆盖 MyBinaryQuantizer 中定义的现有scaling_init：
small_scale_binary_identity = QuantIdentity(
act_quant=MySignedBinaryActQuantizer, scaling_init=0.001, return_quant_tensor=True)
small_scale_binary_identity(torch.randn(4, 4))
QuantTensor(value=tensor([[ 0.0010, 0.0010, 0.0010, -0.0010],
[ 0.0010, -0.0010, 0.0010, -0.0010],
[-0.0010, -0.0010, -0.0010, -0.0010],
[ 0.0010, 0.0010, 0.0010, 0.0010]], grad_fn=<MulBackward0>), scale=tensor(0.0010, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
使用权重统计初始化自定义量化
假设我们想要定义一个二进制权重量化器，其中scaling_impl仍然是ParameterScaling。 然而，我们希望scaling_init 不是用户定义的，而是量化层权重张量中找到的最大值。 为了支持量化器依赖于层的这种用例，量化层会自动将其自身传递给模块名称下的所有量化器。 只需几行代码，我们就可以实现我们的目标：
from brevitas.inject import value
class ParamFromMaxWeightQuantizer(MySignedBinaryWeightQuantizer):
@value
def scaling_init(module):
return module.weight.abs().max()
请注意我们如何利用 @value 装饰器来定义在依赖注入 (DI) 时执行的函数。 这种行为在本质上类似于定义 @property 而不是属性，区别在于 @value 函数可以依赖于注入器的其他属性，这些属性在 DI 期间自动作为函数的参数传入。
将量化器传递给 QuantConv2d 并检索其量化权重：
param_from_max_quant_conv = QuantConv2d(3, 2, (3, 3), weight_quant=ParamFromMaxWeightQuantizer)
param_from_max_quant_conv.quant_weight()
QuantTensor(value=tensor([[[[-0.1876, -0.1876, -0.1876],
[ 0.1876, 0.1876, 0.1876],
[-0.1876, -0.1876, 0.1876]],
[[-0.1876, -0.1876, 0.1876],
[ 0.1876, 0.1876, -0.1876],
[-0.1876, 0.1876, 0.1876]],
[[-0.1876, -0.1876, -0.1876],
[ 0.1876, 0.1876, 0.1876],
[-0.1876, 0.1876, -0.1876]]],
[[[-0.1876, -0.1876, -0.1876],
[ 0.1876, 0.1876, -0.1876],
[ 0.1876, -0.1876, -0.1876]],
[[-0.1876, 0.1876, -0.1876],
[ 0.1876, -0.1876, -0.1876],
[-0.1876, -0.1876, 0.1876]],
[[-0.1876, 0.1876, 0.1876],
[ 0.1876, -0.1876, 0.1876],
[-0.1876, -0.1876, -0.1876]]]], grad_fn=<MulBackward0>), scale=tensor(0.1876, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
假设我们想要在量化模型之上加载一个预训练的浮点权重张量。 我们通过定义具有相同权重形状的单独 nn.Conv2d 层来模拟这种情况：
from torch import nn
float_conv = nn.Conv2d(3, 2, (3, 3))
float_conv.weight.abs().max()
tensor(0.1897, grad_fn=<MaxBackward1>)
然后我们将其加载到 param_from_max_quant_conv 之上：
param_from_max_quant_conv.load_state_dict(float_conv.state_dict())
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-22-5b3646241211> in <module>
----> 1 param_from_max_quant_conv.load_state_dict(float_conv.state_dict())
C:\ProgramData\Miniconda3\envs\pytorch\lib\site-packages\torch\nn\modules\module.py in load_state_dict(self, state_dict, strict)
1405 if len(error_msgs) > 0:
1406 raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-> 1407 self.__class__.__name__, "\n\t".join(error_msgs)))
1408 return _IncompatibleKeys(missing_keys, unexpected_keys)
1409
RuntimeError: Error(s) in loading state_dict for QuantConv2d:
Missing key(s) in state_dict: "weight_quant.tensor_quant.scaling_impl.value".
收到一个错误。 这是因为 ParameterScaling 包含一个学习的 torch.nn.Parameter，并且 Pytorch 期望模型的所有学习参数都包含在正在加载的 state_dict 中。 我们可以通过在 Brevitas 中设置 IGNORE_MISSING_KEYS 配置标志或将 strict=False 传递给 load_state_dict 来解决该问题。 我们选择前者，因为设置 strict=False 对其他类型的问题过于宽容：
from brevitas import config
config.IGNORE_MISSING_KEYS = True
param_from_max_quant_conv.load_state_dict(float_conv.state_dict())
我们也可以通过设置 env 变量来实现相同的目标：BREVITAS_IGNORE_MISSING_KEYS=1
再次看一下量化后的权重：
param_from_max_quant_conv.quant_weight()
QuantTensor(value=tensor([[[[ 0.1897, -0.1897, 0.1897],
[-0.1897, 0.1897, -0.1897],
[-0.1897, 0.1897, -0.1897]],
[[-0.1897, 0.1897, 0.1897],
[ 0.1897, -0.1897, -0.1897],
[ 0.1897, -0.1897, 0.1897]],
[[-0.1897, 0.1897, -0.1897],
[-0.1897, 0.1897, 0.1897],
[-0.1897, 0.1897, 0.1897]]],
[[[ 0.1897, 0.1897, 0.1897],
[-0.1897, 0.1897, -0.1897],
[ 0.1897, 0.1897, -0.1897]],
[[ 0.1897, -0.1897, -0.1897],
[ 0.1897, 0.1897, -0.1897],
[ 0.1897, 0.1897, 0.1897]],
[[-0.1897, 0.1897, -0.1897],
[-0.1897, 0.1897, -0.1897],
[ 0.1897, 0.1897, 0.1897]]]], grad_fn=<MulBackward0>), scale=tensor(0.1897, grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
正如预期的那样，比例因子已更新为新的weight.abs().max()。
内部发生的情况是，在层上调用 load_state_dict之后，再次调用 ParamFromMaxWeightQuantizer.tensor_quant 来重新初始化 BinaryQuant，然后使用基于更新的 module.weight 张量计算的新的scaling_init 值重新初始化 ParameterScaling。 如果没有 ExtendedInjector 的支持，这整个过程就不可能实现。
3. 量化器的共享
简单的将MySignedBinaryWeightQuantizer传递给不同的层，其作用是在不同层之间共享相同的量化策略。每层仍然有自己的量化实现实例：
quant_conv1 = QuantConv2d(3, 2, (3, 3), weight_quant=MySignedBinaryWeightQuantizer)
quant_conv2 = QuantConv2d(3, 2, (3, 3), weight_quant=MySignedBinaryWeightQuantizer)
quant_conv1.weight_quant is quant_conv2.weight_quant
False
共享代理
Beevitas允许在多个层之间共享相同的量化实例。这是通过简单地共享包装它的代理来完成的。这在某些场景中非常有用，例如，我们希望不同的层共享相同的比例因子。语法如下：
quant_conv1 = QuantConv2d(3, 2, (3, 3), weight_quant=MySignedBinaryWeightQuantizer)
quant_conv2 = QuantConv2d(3, 2, (3, 3), weight_quant=quant_conv1.weight_quant)
assert quant_conv1.weight_quant is quant_conv2.weight_quant
True
后台发生的情况是，权重量化器现在可以访问quant_conv1和quant_conv2。 假设我们想要构建一个类似于 ParamFromMaxWeightQuantizer 的量化器，但在这种情况下，我们希望使用两个权重张量的平均值来初始化比例因子。 当量化器可以访问多个父模块时，它们会在依赖项注入时作为元组以与以前相同的名称模块传递。 所以我们可以执行以下操作：
class SharedParamFromMeanWeightQuantizer(MySignedBinaryWeightQuantizer):
@value
def scaling_init(module):
if isinstance(module, tuple):
return torch.cat((module[0].weight.view(-1), module[1].weight.view(-1))).abs().mean()
else:
return module.weight.abs().mean()
quant_conv1 = QuantConv2d(3, 2, (3, 3), weight_quant=SharedParamFromMeanWeightQuantizer)
old_quant_conv1_scale = quant_conv1.quant_weight_scale()
quant_conv2 = QuantConv2d(3, 2, (3, 3), weight_quant=quant_conv1.weight_quant)
new_quant_conv1_scale = quant_conv1.quant_weight_scale()
assert not (old_quant_conv1_scale == new_quant_conv1_scale).item()
False
当使用 quant_conv1 的weight_quant 初始化 quant_conv2 时，两个层的权重量化都会重新初始化，以便它们最终具有相同的比例。
We can see in this example how Brevitas works consistently with Pytorch’s eager execution model. When we initialize quant_conv1 we still don’t know that its weight quantizer is going to be shared with quant_conv2, and the semantics of Pytorch impose that quant_conv1 should work correctly both before and after quant_conv2 is declared. The way we take advantage of dependency injection allows to do so.
共享激活量化实例
共享激活量化的实例更容易，因为对于大多数情况，只需共享整个层本身就足够了，例如 在前向传递中从多个位置调用相同的 QuantReLU。
对于那些无法共享整个层的场景，需要记住一些重要的事情。 激活量化的实例包括（出于性能原因）非线性激活本身的实现（如果有）。因此，例如，应避免使用 QuantReLU.act_quant 初始化 QuantConv2d.output_quant，因为我们不仅不会共享量化器，而且还不会共享 relu 激活函数。
一般来说，激活量化实例的共享应该仅在同类激活之间进行(好像是输出激活只能共享给输出激活)。
4. 权重初始化
有一种情况 Brevitas 无法自动处理。 也就是说，量化器的初始化取决于其所应用的层（例如 ParamFromMaxWeightQuantizer 或 SharedParamFromMeanWeightQuantizer 量化器），但层在初始化后会被修改。
典型的例子是从头开始训练时进行权重初始化（而不是从浮点 state_dict 加载）：
quant_conv_w_init = QuantConv2d(3, 2, (3, 3), weight_quant=ParamFromMaxWeightQuantizer)
torch.nn.init.uniform_(quant_conv_w_init.weight)
assert not (quant_conv_w_init.weight.abs().max() == quant_conv_w_init.quant_weight_scale()).item()
比例因子不再正确初始化。 在这种情况下，我们可以简单地手动触发权重量化器的重新初始化：
quant_conv_w_init.weight_quant.init_tensor_quant()
assert (quant_conv_w_init.weight.abs().max() == quant_conv_w_init.quant_weight_scale()).item()
5. 构建自定义量化 API
假设我们想要分别为权重和激活构建两个量化器，并在它们之上构建一个简单的 API。特别是，我们希望能够在BinaryQuant和ClampedBinaryQuant（带钳位的二进制量化的一种变体）之间切换，并且我们希望有选择地执行每通道缩放。为此，我们将通过 ExtendedInjector 的层次结构来实现控制逻辑，留下两个布尔标志作为量化器的参数公开，然后可以通过相应量化层的关键字参数来设置标志。
from brevitas.core.quant import ClampedBinaryQuant
from brevitas.proxy import WeightQuantProxyFromInjector, ActQuantProxyFromInjector
from brevitas.inject import this
class CommonQuantizer(ExtendedInjector):
scaling_impl = ParameterScaling
signed=True
@value
def tensor_quant(is_clamped):
# returning a class to auto-wire from a value function
# wouldn't be allowed in a standard Injector
if is_clamped:
return ClampedBinaryQuant
else:
return BinaryQuant
@value
def scaling_shape(scaling_per_output_channel):
if scaling_per_output_channel:
# returning this.something from a value function
# wouldn't be allowed in a standard Injector
return this.per_channel_broadcastable_shape
else:
return ()
class AdvancedWeightQuantizer(CommonQuantizer):
proxy_class = WeightQuantProxyFromInjector
@value
def per_channel_broadcastable_shape(module):
return (module.weight.shape[0], 1, 1, 1)
@value
def scaling_init(module, scaling_per_output_channel):
if scaling_per_output_channel:
num_ch = module.weight.shape[0]
return module.weight.abs().view(num_ch, -1).max(dim=1)[0].view(-1, 1, 1, 1)
else:
return module.weight.abs().max()
class AdvancedActQuantizer(CommonQuantizer):
scaling_init = 0.01
proxy_class = ActQuantProxyFromInjector
第一个是 @value 函数可以返回一个类来自动装配和注入，如 tensor_quant 的定义所示。 对于标准注入器来说，这通常是不可能的，但对于扩展注入器来说是可能的。 这样我们就可以在tensor_quant的不同实现之间进行切换。
第二个是特殊对象 this。 它已经存在于依赖项库中，并且用作从量化器本身检索量化器属性的方法。 但是，通常不可能从 @value 函数返回对此的引用。 同样，这是只有 ExtendedInjector 支持的东西，并且它允许以某种方式链接不同的属性，以便仅在必要时计算链接的值。
per_channel_quant_conv = QuantConv2d(
3, 2, (3, 3),
weight_quant=AdvancedWeightQuantizer,
weight_is_clamped=False,
weight_scaling_per_output_channel=True)
per_channel_quant_conv.quant_weight()
QuantTensor(value=tensor([[[[-0.1842, 0.1842, -0.1842],
[-0.1842, -0.1842, 0.1842],
[-0.1842, -0.1842, 0.1842]],
[[-0.1842, -0.1842, 0.1842],
[ 0.1842, -0.1842, 0.1842],
[ 0.1842, 0.1842, -0.1842]],
[[-0.1842, -0.1842, 0.1842],
[ 0.1842, 0.1842, 0.1842],
[-0.1842, 0.1842, -0.1842]]],
[[[ 0.1838, 0.1838, 0.1838],
[-0.1838, -0.1838, -0.1838],
[ 0.1838, 0.1838, -0.1838]],
[[ 0.1838, -0.1838, 0.1838],
[ 0.1838, 0.1838, 0.1838],
[-0.1838, 0.1838, -0.1838]],
[[-0.1838, 0.1838, -0.1838],
[ 0.1838, -0.1838, -0.1838],
[ 0.1838, -0.1838, 0.1838]]]], grad_fn=<MulBackward0>), scale=tensor([[[[0.1842]]],
[[[0.1838]]]], grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
可以加载之前定义的浮点状态字典并观察它如何触发权重标度的更新：
per_channel_quant_conv.load_state_dict(float_conv.state_dict())
per_channel_quant_conv.quant_weight()
QuantTensor(value=tensor([[[[ 0.1875, -0.1875, 0.1875],
[-0.1875, 0.1875, -0.1875],
[-0.1875, 0.1875, -0.1875]],
[[-0.1875, 0.1875, 0.1875],
[ 0.1875, -0.1875, -0.1875],
[ 0.1875, -0.1875, 0.1875]],
[[-0.1875, 0.1875, -0.1875],
[-0.1875, 0.1875, 0.1875],
[-0.1875, 0.1875, 0.1875]]],
[[[ 0.1897, 0.1897, 0.1897],
[-0.1897, 0.1897, -0.1897],
[ 0.1897, 0.1897, -0.1897]],
[[ 0.1897, -0.1897, -0.1897],
[ 0.1897, 0.1897, -0.1897],
[ 0.1897, 0.1897, 0.1897]],
[[-0.1897, 0.1897, -0.1897],
[-0.1897, 0.1897, -0.1897],
[ 0.1897, 0.1897, 0.1897]]]], grad_fn=<MulBackward0>), scale=tensor([[[[0.1875]]],
[[[0.1897]]]], grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
本例中，我们有一个每通道量化器，因此原始浮点权重张量现在按通道进行量化。
同样，可以将自定义激活量化器应用于QuantIdentity层：
from brevitas.nn import QuantIdentity
quant_identity = QuantIdentity(
act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=False)
quant_identity(torch.randn(4, 4))
tensor([[-0.0100, -0.0100, 0.0100, -0.0100],
[-0.0100, -0.0100, -0.0100, 0.0100],
[-0.0100, 0.0100, 0.0100, 0.0100],
[-0.0100, 0.0100, 0.0100, 0.0100]], grad_fn=<MulBackward0>)
注意 AdvancedActQuantizer 没有定义 per_channel_broadcastable_shape，但不会触发任何错误。这是因为仅当scaling_per_output_channel为True时才需要this.per_channel_broadcastable_shape。将scaling_per_output_channel设置为True：
from brevitas.nn import QuantIdentity
quant_identity = QuantIdentity(
act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=True)
---------------------------------------------------------------------------
DependencyError Traceback (most recent call last)
<ipython-input-36-b3479e90d1a9> in <module>
2
3 quant_identity = QuantIdentity(
----> 4 act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=True)
c:\brevitas_fx\src\brevitas\nn\quant_activation.py in __init__(self, act_quant, return_quant_tensor, **kwargs)
134 act_quant=act_quant,
135 return_quant_tensor=return_quant_tensor,
--> 136 **kwargs)
137
c:\brevitas_fx\src\brevitas\nn\quant_layer.py in __init__(self, act_impl, passthrough_act, input_quant, act_quant, return_quant_tensor, **kwargs)
77 passthrough_act,
78 act_quant,
---> 79 **kwargs)
80
81 @property
c:\brevitas_fx\src\brevitas\nn\mixin\act.py in __init__(self, act_impl, passthrough_act, act_quant, **kwargs)
157 proxy_prefix='act_',
158 kwargs_prefix='',
--> 159 **kwargs)
160
161 @property
c:\brevitas_fx\src\brevitas\nn\mixin\base.py in __init__(self, quant, proxy_protocol, none_quant_injector, proxy_prefix, kwargs_prefix, **kwargs)
98 quant_injector = quant
99 quant_injector = quant_injector.let(**filter_kwargs(kwargs_prefix, kwargs))
--> 100 quant = quant_injector.proxy_class(self, quant_injector)
101 else:
102 if not isinstance(quant, proxy_protocol):
c:\brevitas_fx\src\brevitas\proxy\runtime_quant.py in __init__(self, quant_layer, quant_injector)
108
109 def __init__(self, quant_layer, quant_injector):
--> 110 super(ActQuantProxyFromInjector, self).__init__(quant_layer, quant_injector)
111 self.is_passthrough_act = _is_passthrough_act(quant_injector)
112
c:\brevitas_fx\src\brevitas\proxy\quant_proxy.py in __init__(self, quant_layer, quant_injector, export_mode, export_handler)
74 # Use a normal list and not a ModuleList since this is a pointer to parent modules
75 self.tracked_module_list = []
---> 76 self.add_tracked_module(quant_layer)
77 self.export_handler = export_handler
78 self.export_mode = export_mode
c:\brevitas_fx\src\brevitas\proxy\quant_proxy.py in add_tracked_module(self, module)
130 self.tracked_module_list.append(module)
131 self.update_tracked_modules()
--> 132 self.init_tensor_quant()
133 else:
134 raise RuntimeError("Trying to add None as a parent module.")
c:\brevitas_fx\src\brevitas\proxy\runtime_quant.py in init_tensor_quant(self)
120
121 def init_tensor_quant(self):
--> 122 tensor_quant = self.quant_injector.tensor_quant
123 act_impl = self.quant_injector.act_impl
124 is_act_enabled = _is_act_enabled(act_impl, tensor_quant)
[... skipping hidden 1 frame]
C:\ProgramData\Miniconda3\envs\pytorch\lib\site-packages\_dependencies\this.py in __call__(self, __self__)
49 if kind == ".":
50 try:
---> 51 result = getattr(result, symbol)
52 except DependencyError:
53 message = (
[... skipping hidden 1 frame]
DependencyError: 'AdvancedActQuantizer' can not resolve attribute 'per_channel_broadcastable_shape'
正如预期的那样，我们收到一个错误，指出量化器无法解析per_channel_broadcastable_shape。如果我们将其传入，那么我们可以获得每通道量化器：
quant_identity = QuantIdentity(
act_quant=AdvancedActQuantizer, is_clamped=True, scaling_per_output_channel=True,
per_channel_broadcastable_shape=(4, 1), return_quant_tensor=True)
quant_identity(torch.randn(4, 4))
QuantTensor(value=tensor([[-0.0100, 0.0100, -0.0100, -0.0100],
[-0.0100, 0.0100, -0.0100, -0.0100],
[ 0.0100, -0.0100, 0.0100, -0.0100],
[ 0.0100, -0.0100, -0.0100, -0.0100]], grad_fn=<MulBackward0>), scale=tensor([[0.0100],
[0.0100],
[0.0100],
[0.0100]], grad_fn=<AbsBinarySignGradFnBackward>), zero_point=tensor(0.), bit_width=tensor(1.), signed_t=tensor(True), training_t=tensor(True))
我们已经看到依赖注入有多么强大。 从某种意义上说，它甚至过于富有表现力。 对于对构建完全自定义量化器不感兴趣的用户来说，可能很难理解如何根据最佳实践将 brevitas.core 下可用的各种组件组装在一起。