深度学习框架量化感知训练的思考及OneFlow的一种解决方案
Posted just_sort
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了深度学习框架量化感知训练的思考及OneFlow的一种解决方案相关的知识，希望对你有一定的参考价值。
【GiantPandaCV导语】这篇文章分享的是笔者最近在OneFlow做的一个项目，将Pytorch FX移植到OneFlow之后实现了自动量化感知训练动态图模型(在Pytorch和OneFlow中都称为nn.Module)。现在用户可以在自己构建的nn.Module基础上，修改很少的代码即可完成从nn.Module量化感知训练到用TensorRT将量化感知训练后的模型部署到GPU上运行的完整链路。在TensorRT上推理是利用了ONNX作为中间表示，即Oneflow动态图模型(nn.Module)->OneFlow量化感知训练模型(nn.Module)->OneFlow静态图(nn.Graph)->ONNX->TensorRT。量化感知训练是基于支持在Eager下写Pass的FX模块（FX被Pytorch率先提出，笔者将其基础设施移植到了OneFlow）来完成的。读者如果想体验这个功能可以按照本文的方法进行操作，有任何使用上的问题可以联系笔者。
0x0. 总览
好久不见，大家国庆快乐！
相信不少小伙伴都了解或者使用了一些深度学习框架比如Pytorch，TensorFlow，OneFlow（也是笔者目前正在参与开发的）。但当大家使用深度学习框架的训练量化方案时如果第一感觉就是太复杂了，那么你可能会对这篇文章感兴趣！因为我在2个月前开始接触这个项目前，对量化感知训练的知识积累也非常少，并且我也会认为各个框架的量化感知训练方案很复杂，甚至不想研究这些API。
这篇文章我会以Pytorch的两代量化方案开始切入谈一谈他们的好处和坏处，然后我会讲讲我在吸收了Pytorch的部分优秀成果(FX模块)并加上一些自己的想法后把OneFlow的量化感知训练方案做成了什么样子。这里先罗列一下这篇文章中涉及到的知识点：
- Pytorch FX模块
- Eager Pass
- 量化感知训练
- Conv+BN的融合
- OneFlow的动静转换(nn.Graph)
- ONNX
- TensorRT
如果你对上面的任意一个知识点不熟悉，那也是完全没有关系的。实际上即使你只会用Pytorch搭建模型也可以快速把本文的量化感知训练方案用起来。因为量化感知训练的工作和模型转化成ONNX以及用TensorRT来部署运行的代码我们在OneFlow社区中均开源了。
简单总结一下就是，用户可以基于OneFlow搭建一个动态图模型（即nn.Module，算子的API和Pytorch基本一样），然后调用下面的几行代码就可以完成这个动态图模型（是一个nn.Module）自动在合适的位置插入量化模块生成一个量化模型（仍然是nn.Module），然后基于这个量化模型完成量化感知训练。
gm: flow.fx.GraphModule = flow.fx.symbolic_trace(net)
qconfig = {
'quantization_bit': 8,
'quantization_scheme': "symmetric",
'quantization_formula': "cambricon",
'per_layer_quantization': True,
'momentum': 0.95,
}
net = quantization_aware_training(gm, flow.randn(1, 3, 32, 32), qconfig)
net = net.to(device)
在训练完成后，调用下面的代码完成训练量化模型到ONNX的转换，并使用TensorRT在GPU上推理。
quantization_resnet18 = quantization_aware_training(gm, flow.randn(1, 3, 32, 32).to("cuda"), qconfig)
quantization_resnet18 = quantization_resnet18.to("cuda")
quantization_resnet18.eval()
checkpoint = flow.load('/home/zhangxiaoyu/oneflow-cifar/checkpoint/epoch_11_val_acc_83.280000')
quantization_resnet18.load_state_dict(checkpoint)
origin_gm: flow.fx.GraphModule = flow.fx.symbolic_trace(resnet18)
dequantization_resnet18 = dequantization_aware_training(origin_gm, gm, flow.randn(1, 3, 32, 32).to("cuda"), qconfig)
dequantization_resnet18 = dequantization_resnet18.to("cuda")
dequantization_resnet18.eval()
class ResNet18Graph(flow.nn.Graph):
def __init__(self):
super().__init__()
self.m = dequantization_resnet18
def build(self, x):
out = self.m(x)
return out
def test_resnet():
resnet_graph = ResNet18Graph()
resnet_graph._compile(flow.randn(1, 3, 32, 32).to("cuda"))
with tempfile.TemporaryDirectory() as tmpdirname:
flow.save(dequantization_resnet18.state_dict(), tmpdirname)
convert_to_onnx_and_check(resnet_graph, flow_weight_dir=tmpdirname, onnx_model_path="/tmp", print_outlier=True)
ipt_dict, onnx_res = run_onnx("/tmp/model.onnx", get_onnx_provider("cpu"))
trt_res = run_tensorrt("/tmp/model.onnx", ipt_dict[list(ipt_dict.keys())[0]])
compare_result(onnx_res, trt_res, atol=1e-4, print_outlier=True)
test_resnet()
用户只需要使用上面示例中的短短几十行代码就可以完成一个端到端的量化感知训练到GPU部署的全流程。所以我认为这项工作是有趣并且相对简洁的，当然我更希望听到用户的想法，然后就写了这篇文章来分享这个项目。这个项目的所有代码均开源在了OneFlow社区，下面是对应的链接。如果你使用这个方案碰到了什么问题都可以第一时间联系我。我的个人微信号是bbuf23333，来时请备注 量化感知训练
- OneFlow FX（用来实现量化感知训练的基础设施）：https://github.com/Oneflow-Inc/oneflow/pull/5939
- OneFlow Cifar（基于OneFlow FX量化训练Cifar10）：https://github.com/BBuf/oneflow-cifar
- OneFlow->ONNX和TensorRT运行：https://github.com/Oneflow-Inc/oneflow_convert/pull/45
0x1. Pytorch量化方案的沉浮
这一节主要基于Pytorch的官方文档：https://pytorch.org/docs/1.9.0/quantization.html来进行说明。Pytorch第一代量化方案叫作Eager Mode Quantization，然后从1.8开始推出FX Graph Mode Quantization。Eager Mode Quantization需要用户手动更改模型，并手动指定需要融合的Op。FX Graph Mode Quantization解放了用户，一键自动量化，无需用户修改模型和关心内部操作。这个改动具体可以体现在下面的图中。
下面分别解释一下Pytorch这两种量化方式的区别。
Eager Mode Quantization
class Net(nn.Module):
def __init__(self, num_channels=1):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(num_channels, 40, 3, 1)
self.conv2 = nn.Conv2d(40, 40, 3, 1)
self.fc = nn.Linear(5*5*40, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2, 2)
x = x.reshape(-1, 5*5*40)
x = self.fc(x)
return x
Pytorch可以在nn.Module的foward里面随意构造网络，可以调用其它nn.Module，也可以调用nn.functional.xxx，甚至可以在里面写If这种控制逻辑。但这也带来了一个问题，就是在Eager层面比较难获取这个模型的图结构。所以在Eager Mode Quantization中，要量化这个网络必须做手动修改：
class NetQuant(nn.Module):
def __init__(self, num_channels=1):
super(NetQuant, self).__init__()
self.conv1 = nn.Conv2d(num_channels, 40, 3, 1)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(40, 40, 3, 1)
self.relu2 = nn.ReLU()
self.pool2 = nn.MaxPool2d(2, 2)
self.fc = nn.Linear(5*5*40, 10)
self.quant = torch.quantization.QuantStub()
self.dequant = torch.quantization.DeQuantStub()
def forward(self, x):
x = self.quant(x)
x = self.relu1(self.conv1(x))
x = self.pool1(x)
x = self.relu2(self.conv2(x))
x = self.pool2(x)
x = x.reshape(-1, 5*5*40)
x = self.fc(x)
x = self.dequant(x)
return x
也就是说，除了Conv，Linear这些含有参数的Module外，ReLU，MaxPool2d也要在__init__中定义，Eager Mode Quantization才可以正确处理。
除了这一点，还有一些情况是需要Fuse之后做量化比如Conv+ReLU，那么还需要手动指定这些层进行折叠，目前这种量化模式支持ConV + BN、ConV + BN + ReLU、Conv + ReLU、Linear + ReLU、BN + ReLU的折叠。
model = NetQuant()model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
modules_to_fuse = [['conv1', 'relu1'], ['conv2', 'relu2']] # 指定合并layer的名字
model_fused = torch.quantization.fuse_modules(model, modules_to_fuse)
model_prepared = torch.quantization.prepare(model_fused)
post_training_quantize(model_prepared, train_loader) # 这一步是做后训练量化
model_int8 = torch.quantization.convert(model_prepared)
整个流程比较逆天，不知道有没有人用。不过公众号有小伙伴确实用过，见文章：Pytorch量化感知训练详解
FX Graph Mode Quantization
关于Pytorch FX模块是什么，我们放到下一节来讲。
由于 Pytorch FX 可以自动跟踪 forward 里面的代码，因此它是真正记录了网络里面的每个节点，在 fuse 和动态插入量化节点方面，比 Eager 模式强太多。对于前面那个模型代码，我们不需要对网络做修改，直接让 FX 帮我们自动修改网络即可，一个使用示例如下：
from torch.quantization import get_default_qconfig, quantize_jit
from torch.quantization.quantize_fx import prepare_fx, convert_fx
model = Net()
qconfig = get_default_qconfig("fbgemm")
qconfig_dict = {"": qconfig}
model_prepared = prepare_fx(model, qconfig_dict)
post_training_quantize(model_prepared, train_loader) # 这一步是做后训练量化
model_int8 = convert_fx(model_prepared)
基于这两套量化方案来看，基于FX的量化方案显然更加优秀，因为它不需要用户在定义模型的时候做什么额外限制，用户仍然是随心所欲的写模型代码就行了，这才符合人的常识。我在做OneFlow的量化感知训练方案时也正是基于FX这个基础设施（我将其核心功能移植到了OneFlow框架下，代码链接第一节给了）来完成的。
另外在TensorRT的工程中：https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization发现Pytorch量化模型要转为ONNX来部署现在似乎还是得基于第一个版本的方案，Pytorch FX这边似乎想直接从nn.Module转到TensorRT，不经过ONNX的中间表示，所以我这里的技术路线还是有点不一样。
0x2. OneFlow FX (在Eager中写Pass)
FX可以用来做什么？
FX可以将一个nn.Module变换后生成另外一个nn.Module，只需要在这个架构的基础上实现一些Transformation（也可以叫Pass），比如在Conv后自动插入伪量化节点实现训练量化，然后生成GraphModule（这个也是nn.Module）进行训练和转为ONNX进行部署。
OneFlow FX模块在这个PR（https://github.com/Oneflow-Inc/oneflow/pull/5939）中实现，这里复用了Pytorch FX基础设施的核心逻辑和代码，这个PR里的主要工作为：
- 精简Pytorch FX的特殊设计比如对_C的Trace，和Jit的交互。保留核心功能，即Symbolic Tracing，Intermediate Representation和Transformation以及Python Codegen这4个组成部分。
- 分步实现以上四大功能的代码，完全适配OneFlow的相关设计，现在可以一键import oneflow.fx来体验。可以Trace住基本所有OneFlow API搭建的Eager模型的结构，并将其变换成一个等价的
nn.Module，我们还可以在这个nn.Module的基础上自定义自己的Transformation Pass，我这里实现了Shape Infer和Quantization以及Dequantization的Pass。 - 增加AlexNet，ResNet50，MobileNetV2等模型的测试。
然后分享一下OneFlow FX的整体思路。
先看一个示例：
import oneflow
# Simple module for demonstration
class MyModule(oneflow.nn.Module):
def __init__(self):
super().__init__()
self.param = oneflow.nn.Parameter(oneflow.rand(3, 4))
self.linear = oneflow.nn.Linear(4, 5)
def forward(self, x):
return self.linear(x + self.param).clamp(min=0.0, max=1.0)
module = MyModule()
from oneflow.fx import symbolic_trace
# Symbolic tracing frontend - captures the semantics of the module
symbolic_traced : oneflow.fx.GraphModule = symbolic_trace(module)
# High-level intermediate representation (IR) - Graph representation
print(symbolic_traced.graph)
"""
graph():
%x : [#users=1] = placeholder[target=x]
%param : [#users=1] = get_attr[target=param]
%add : [#users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {})
%linear : [#users=1] = call_module[target=linear](args = (%add,), kwargs = {})
%clamp : [#users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0})
return clamp
"""
# Code generation - valid Python code
print(symbolic_traced.code)
"""
def forward(self, x):
param = self.param
add = x + param; x = param = None
linear = self.linear(add); add = None
clamp = linear.clamp(min = 0.0, max = 1.0); linear = None
return clamp
"""
在FX中有一个Proxy类，它会把oneflow中所有的call_method和call_function以及math库中的函数和常见的魔法函数都包装一遍来记录OneFlow中所有的运算符，这个在import oneflow.fx时就做好了。然后在传入一个nn.Module调用symbolic_trace进行跟踪代码的时候会首先处理__init__中的其它nn.Module，把这些nn.Module也用Proxy包起来，同时输入数据也要包起来。
用Proxy包好所有程序中可能存在的运算符之后就执行一遍forward，这个forward的输入数据不再是Tensor而是Proxy(Tensor)。由于程序的执行过程类似于一个运算符和数据入栈出栈的过程，所以我们可以直接按照这个执行顺序将刚才用Proxy记录下来的数据和Op进行unpack，unpack之后可以拿到真实的Tensor, Parameter和运算符等等，我们将这些数据和运算符当作点和边去构造一个新的Graph。那么Graph是怎么转化成nn.Module的呢？FX中通过引入GraphModule的数据结构来持有这个Graph，此外GraphModule还持有code和foward成员，这两者都是基于Graph自动生成的，注意GraphModule仍然是nn.Module。
自动生成的代码就是GraphModule中的code，打印出来其实就是整个forward函数的完整执行过程。
另外FX还提供了一个Interpreter类用来让用户自定义nn.Module的执行过程，比如这个PR提供了一个基于这个类做所有中间Tensor形状推导的Pass。另外还提供了一个基于pydot将GraphModule结构可视化的Pass，如下图。
相信到这里大家对FX有一个了解了，这里最棒的一个功能就是我们可以对nn.Module进行修改，然后返回变化后的nn.Module。说
到这里，我们自然能想到量化感知训练不就是把Conv+BN或者Conv，Linear等组件替换为插入了伪量化节点的组件吗？所以我们基于FX来写一个Pass就可以完成这件事了。
这就是上面说的，FX支持在Eager写Pass。
然而FX也存在缺陷，目前无法处理控制流，需要注意网络中不要带控制流（不过这一点暂时影响不大，因为用户一般都不会部署含有控制流的网络，如果有这个需求我们可以交流）。
0x3. 实现量化感知训练Pass
有了OneFlow FX之后我们就可以实现一个量化感知训练的Pass来将用户自定义的网络中自动插入量化感知训练组件来完成量化感知训练了。
以ResNet18为例，它只有Conv+BN这种模式，即对于任意一个卷积层后面都跟了一个BN层，在推理的时候TensorRT会做Conv+BN的融合，那么我们在训练的时候也是必须要做Conv+BN的融合的，不然会影响部署的精度。所以，我们首先需要把BN层的参数和卷积层的参数融合，然后再对这个参数做量化，具体过程如下图所示：
下面是Conv和BN融合的公式：
y b n = γ ( y − μ σ ) + β y_{bn}=\\gamma(\\frac{y-\\mu}{\\sigma})+\\beta ybn=γ(σy−μ)+β
= γ ( W x + b − μ σ ) + β =\\gamma(\\frac{Wx+b-\\mu}{\\sigma})+\\beta =γ(σWx+b−μ)+β
= γ W x σ + γ b − μ σ + β =\\gamma\\frac{Wx}{\\sigma}+\\gamma\\frac{b-\\mu}{\\sigma}+\\beta =γσWx+γσb−μ+β
所以：
W m e r g e = γ W σ W_{merge}=\\gamma\\frac{W}{\\sigma} Wmerge=γσW
b m e r g e = γ b − μ σ + β b_{merge}=\\gamma\\frac{b-\\mu}{\\sigma}+\\beta bmerge=γσb−μ+β
公式中的， W W W和 b b b分别表示卷积层的权值与偏置， x x x和 y y y分别为卷积层的输入与输出，则根据 B N BN BN的计算公式，可以推出融合了batchnorm参数之后的权值与偏置， W m e r g e W_{merge} W基于OneFlow实现量化感知训练
适配PyTorch FX，OneFlow让量化感知训练更简单