Pytorch推出fx，量化起飞

(本文首发于公众号，没事来逛逛)

Pytorch1.8 发布后，官方推出一个 torch.fx 的工具包，可以动态地对 forward 流程进行跟踪，并构建出模型的图结构。这个新特性能带来什么功能呢？别的不说，就模型量化这一块，炼丹师们有福了。

其实早在三年前 pytorch1.3 发布的时候，官方就推出了量化功能。但我觉得当时官方重点是在后端的量化推理引擎（FBGEMM 和 QNNPACK）上，对于 pytorch 前端的接口设计很粗糙。用过 pytorch 量化的同学都知道，这个量化接口实在是太麻烦、太粗糙、太暴力了。官方又把这个第一代的量化方式称为 Eager Mode Quantization。我后面会用一个例子来展示这种方式有多傻x。

而随着 fx 的推出，由于可以动态地 trace 出网络的图结构，因此就可以针对网络模型动态地添加一些量化节点。官方又称这种新的量化方式为 FX Graph Mode Quantization。上一张官方的图来对比一下这两种方式的优缺点：

我总结一下这张图，Eager Mode Quantization 需要手工修改网络代码，并对很多节点进行替换，而 FX Graph 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 这种动态图有一个很好的地方，就是可以在 forward 函数中天马星空构造「电路图」，比如 Functional 这些函数模块可以随意调用，而不需要在 init 函数里面事先定义，再比如可以随时加入 if、for 等逻辑控制语句。这就是动态图区别于静态图的地方。但这种好处的代价就是，我们很难获取网络的图结构。

下面我们就看看 Eager 模式下的量化怎么操作。

看过我之前量化系列教程的读者应该知道，模型量化需要在原网络节点中插入一些伪量化节点，或者把一些 Module 或者 Function 替换成量化的形式。对于 Eager 模式，由于它只会对 init 函数里面定义的模块进行替换，因此，如果有一些 op 没有在 init 中定义，但又在 forward 中用到了（比如上面代码的 F.relu），那就凉凉了。

因此，上面这段网络代码是没法直接用 Eager 模式量化的，需要重新写成下面这种形式：

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 模式才能进行处理。

这还没完，由于有些节点是要做 fuse 之后才能量化的（比如：Conv + 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 这套设计直接劝退了很多想上手量化的同学，我很早之前看到这些操作也是一点上手的欲望都没有。

那这套新的 Graph 模式的量化又是怎样的呢？

由于 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)

对比一下前面 Eager 模式的流程，有没有感觉自己又可以了。

目前 FX 这个新工具包还在优化中，很多功能并不完善。比如，如果 forward 代码中出现了 if 和 for 等控制语句，它依然还是解析不了，这个时候就需要你把 if 还有 for 语句手动拆解掉。但相比起之前的流程，已经是一个巨大的进步了。而且，有了这个图结构，很多后训练量化的算法也可以更加方便的操作（很多 PTQ 的算法需要针对针对网络的拓扑结构优化）。除此以外，像 NAS 等模型结构搜索之类的算法，也可以更加方便的进行。

总的来说，pytorch 推出的这个新特性实在是极大弥补了动态图的先天不足。之前一直考虑针对 pytorch 做一些离线量化的工具，但由于它的图结构很难获取，因此一直难以入手（ONNX 和 jit 这些工具对量化支持又不够）。现在有了 fx，感觉可以加油起飞了。希望官方再接再厉，不要机毁人亡。

欢迎关注我的公众号：大白话AI，立志用大白话讲懂AI。

Pytorch推出fx，量化起飞相关推荐

PyTorch笔记——FX
官方文档链接:https://pytorch.org/docs/master/fx.html# 概述 FX是供开发人员用于转换nn.Module实例的工具包.FX由三个主要组件组成:符号追踪:symb ...
pytorch FX模型静态量化
文章目录前言一.pytorch静态量化(手动版) 踩坑: 二.使用FX量化 1.版本 2.代码如下: 总结前言以前面文章写到的mobilenet图像分类为例,本文主要记录一下pytorchh训 ...
Pytorch模型量化实践并以ResNet18模型量化为例(附代码)
更多.更及时内容欢迎微信公众号:小窗幽记机器学习围观,后续会进一步整理模型推理加速和部署方面的相关内容. 文章目录量化基础知识映射函数量化参数校准(Calibration) Affine和S ...
PyTorch 1.5发布，与AWS联手推出TorchServe
导读:近日 PyTorch 发布了 1.5 版本的更新,作为越来越受欢迎的机器学习框架,PyTorch 本次也带来了大的功能升级.此外, Facebook 和 AWS 还合作推出了两个重要的 PyTo ...
如何使用PyTorch的量化功能？
背景在深度学习中,量化指的是使用更少的 bit 来存储原本以浮点数存储的 tensor,以及使用更少的 bit 来完成原本以浮点数完成的计算.这么做的好处主要有如下几点: 更少的模型体积,接近 4 ...
掌握Pytorch模型压缩裁剪与量化
在深度学习模型的搭建和部署中,我们需要考虑到模型的权重个数.模型权重大小.模型推理速度和计算量.本文将分享在Pytorch中进行模型压缩.裁剪和量化的教程. 权重压缩模型在训练时使用的模型权重类型为 ...
利用torch.fx进行后量化
torch.fx 量化支持--FX GRAPH MODE QUANTIZATION torch.fx目前支持的量化方式: Post Training Quantization Weight Only ...
模型压缩：模型量化打怪升级之路-工具篇
本文转载自商汤泰坦公开课. 1/ 最近发现一些还在学校读书的同学非常关注一个量化工作精度的高低,读过我上篇分享(模型压缩:模型量化打怪升级之路 - 0 序章)的同学应该知道,部分学术界的工作与工业界的 ...
深度学习框架量化感知训练的思考及OneFlow的解决方案
作者 | BBuf 原文首发于公众号GiantPandaCV 0x0.总览相信不少小伙伴都了解或者使用了一些深度学习框架比如PyTorch,TensorFlow,OneFlow(也是笔者目前正在参与 ...

Pytorch推出fx，量化起飞

Pytorch推出fx，量化起飞相关推荐

最新文章

热门文章