torch.fx 简介与量化

        pytroch发布的torch.fx工具包可以说是很好的消除一些动态图和静态图的Gap,可以使得我们对于nn.Module的各种变化操作变得非常简单。

动态图和静态图:

        动态意味着程序将按照我们编写命令的顺序进行执行。这种机制将使得调试更加容易,并且也使得我们将大脑中的想法转化为实际代码变得更加容易。而静态则意味着程序在编译执行时将先生成神经网络的结构,然后再执行相应操作。按照动态图和静态图最明显的两个框架就是pytorch和tensorflow两个。

静态图的优势就在于,因为我们提前获知了模型结构,这就更方便于我们做一些变化,比如量化和算子融合。但是缺点就在于,不够灵活,我们不能随心所欲,天马行空的设计我们的模型结构。而动态图,程序将按照我们编写命令的顺序进行执行。这种机制将使得调试更加容易,并且也使得我们将大脑中的想法转化为实际代码变得更加容易。

        最简单的例子比如,tensorflow中没有forward,而pytorch中有forward.

torch.fx

        torch,fx,这是一个用于捕获和转换PyTorch程序的纯Python系统。主要分为三个结构块:符号追踪器(symbolic tracer),中间表示(intermediate representation), Python代码生成(Python code generation)。

        这三个组件可以做什么?直观看起来,torch.fx做的就是将一个Module转换为静态图。首先,通过追踪器获取到模型的graph,从而产生中间表示,对于中间表示我们做一系列变化,再通过python代码生成来生成python代码。

        这段代码在官网上都有:torch.fx — PyTorch 1.10.1 documentation

        首先,通过symbolic_traced来跟踪模型,得到模型的graph,我们可以对这个graph做出一系列变化,后通过代码生成得到对应的python代码,来修改网络结构。

        这里来看一个修改Module的简单例子,这个例子中我们将模块中所有的操作替换成 torch.mul() :

        

torch.fx与量化

        在torch.fx提出之前,pytorch的量化方式官方称为 Eager Mode Quantization,而随着torch.fx的提出,由于可以动态地 trace 出网络的图结构,因此就可以针对网络模型动态地添加一些量化节点。官方又称这种新的量化方式为 FX Graph Mode Quantization。

        两个方式的区别主要在于,Eager Mode Quantization 需要手工修改网络代码,并对很多节点进行替换,而 FX Graph Mode Quantization 则大大提高了自动化的能力。

        举个例子:

         定义一个网络模型

        在Eager Mode Quantization模式下对模型进行量化:

        这意味着,我们需要去修改网络结构,同时由于有些节点是要做 fuse 之后才能量化的(比如:Conv + ReLU),因此我们需要手动指定这些层进行合并。

        简单的网络还行,一旦复杂,各种天马行空的结构,复杂程度成倍增加。

         在FX Graph Mode Quantization。模式下对模型进行量化:

         非常简单。

        接下来,简单介绍一下代码结构与原理:

        pytorch为我们提供了这几种量化方式:

        PTQ和QAT就不说了。

        模型静态量化:考虑到动态量化的量化方式有时会带来模型预测效果的大幅度下降,因此引入静态量化,它同样发生在模型训练后,为了判断哪些权重或激活值应该被量化,哪些应该保留或小幅度量化,在预测过程开始前,在模型中节点插入“观测者”(衡量节点使用情况的一些计算方法),他们将在一些实验数据中评估节点使用情况,来决定是否将其权重或激活值进行量化,因为在预测过程中,这些节点是否被量化已经确定,因此称作静态量化。

        模型动态量化:操作最简单也是压缩效果最好的量化方式,量化过程发生在模型训练后,针对模型权重采取量化,之后会在模型预测过程中,再决定是否针对激活值采取量化,因此称作动态量化(在预测时可能发生量化)

        模型静态量化可以分为两种:仅weight量化以及weight与activation同时量化。前者主要目的是使得模型参数所占内存减小,且相对简单。后者则相对较为复杂。

         这段是官网给出的例子。

        不管是哪种量化方式,代码都差不多。

        第一步,也就是模型,没啥可说的。

        第二步,qconfig,该步骤目的是设置模型量化的方式,通过插入两个observer来监测activation与weight。同时由于推理平台的不同,pytorch的量化配置也不相同。

        官网上也都有介绍。

        第三步,prepare,将每个可支持量化的模块插入Observer,收集数据并进行量化分析。

        第四步,在激活量化时需要,也就是校准,输入一部分数据来校准激活量化的scale和zeropoint。

        第五步,convert,运用代码转化模型。 

        简单来说:

         模型量化后就是量化模型的推理,与浮点模型的区别主要是要进行量化节点的插入以及op替换。

        量化节点插入:需要在网络的forward里面插入QuantStub与DeQuantSub两个节点。

        op替换:需要将模型中的Add、Concat等操作替换为支持量化的FloatFunctional

如有错误,欢迎各位批评指正!!

    

  • 9
    点赞
  • 20
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值