brevitas + FINN｜端到端神经网络部署

brevitas量化 + FINN部署｜导出IP过程的坑坑坑

题主要做的事，是把一个神经网络部署到Pynq z2上。大概流程分为这么几步：

1. 在训练网络的时候用brevitas的框架，因此需要对torch代码做一点小修改
2. 导出模型weight，在FINN框架下导出IP
3. 放到Pynq-z2上验证

先说结论，brevitas用起来大概没什么问题，FINN在最后一步导出的IP的时候遇到了一点数据流的问题，导致相同的input在多次推断后结果不大一样。题主目前在官方的github上提出了issue，但是还没有得到回复。所以最终是没有用FINN的优化，是把brevitas导出的模型手动用C语言复现，用HLS综合并部署的。
不过毕竟把FINN的流程走完了，还是有一些坑。接下来主要说一下brevitas要怎样才能适配到后续的FINN。

如何用brevitas 导出一个完美适配FINN的模型

官方仓库地址：https://github.com/Xilinx/brevitas
虽然这个库的文档里没有说，但是作者有在issue里提到这个库使用的QAT方法是exactly same as 谷歌的经典全整型量化的文章Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference 。这个方法被tensorflow用在他们的量化库里，所以brevitas其实是torch的一个实现该方法的library。

接下来讲几个踩的坑：

1. 用到的层和定义的层需要完全匹配

因为brevitas在训练时，对模型class中的每个attribute 计算scale，而如果这个层在forward函数里被用到的话，这个层是没有quant_scale这些量化相关的attribute的。因此用到的层和定义的层需要完全匹配，在后续导参数的时候才不会出错。举例:

class Quantized_S2_Block(nn.Module):
  	def __init__(self, in_ch, out_ch):
        super(Quantized_S2_Block, self).__init__()
        self.conv0_d = qnn.QuantConv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=(1, 9), stride=2,
                                       padding=(0, 4), bias=False, weight_bit_width=qw, bias_quant=Int32Bias)
        self.relu0 = qnn.QuantReLU(bit_width=qw, return_quant_tensor=True)
	def forward(x):
		out = self.relu0(x)
		return out

这个例子在训练的时候可能不会报错，但是导出的参数是有问题的，因为con0_d在forward函数中没有被用到，所以没有scale。

2. 全连接层不要用1x1conv替代

题主使用的网络最后一层是全连接，之前个人习惯使用1x1conv来替代全连接，like：

       self.conv_block_p = qnn.QuantConv2d(in_channels=n_mels, out_channels=int(16 * k), kernel_size=(1, 1), stride=1, bias=False, weight_bit_width=qw, bias_quant=Int32Bias)

这样是不行的！！必须定义成nn.Linear 才行！不然还是会报错。

3. 每个Conv后面都要接ReLu。ReLu用几个就要定义几个。

因为每个激活函数都是有自己的scale的，所以必须像官方文档里那样定义好几个Relu。以下代码取自官方教程：

from torch.nn import Module
import torch.nn.functional as F

import brevitas.nn as qnn
from brevitas.quant import Int8Bias as BiasQuant


class QuantWeightActLeNet(Module):
    def __init__(self):
        super(LowPrecisionLeNet, self).__init__()
        self.quant_inp = qnn.QuantIdentity(bit_width=4)
        self.conv1 = qnn.QuantConv2d(3, 6, 5, bias=True, weight_bit_width=4)
        self.relu1 = qnn.QuantReLU(bit_width=4)
        self.conv2 = qnn.QuantConv2d(6, 16, 5, bias=True, weight_bit_width=4)
        self.relu2 = qnn.QuantReLU(bit_width=3)
        self.fc1   = qnn.QuantLinear(16*5*5, 120, bias=True, weight_bit_width=4)
        self.relu3 = qnn.QuantReLU(bit_width=4)
        self.fc2   = qnn.QuantLinear(120, 84, bias=True, weight_bit_width=4)
        self.relu4 = qnn.QuantReLU(bit_width=4)
        self.fc3   = qnn.QuantLinear(84, 10, bias=True)

    def forward(self, x):
        out = self.quant_inp(x)
        out = self.relu1(self.conv1(out))
        out = F.max_pool2d(out, 2)
        out = self.relu2(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = out.reshape(out.shape[0], -1)
        out = self.relu3(self.fc1(out))
        out = self.relu4(self.fc2(out))
        out = self.fc3(out)
        return out

quant_weight_act_lenet = QuantWeightActLeNet()

# ... training ...

就像他们给的代码一样，relu需要用几个就定义几个。而且文档里每一个有参数的后面都接了relu，题主也是按照这样的规范定义网络的。

4. Pooling的问题

一般大家都用avg_pool 或者max_pool， 题主的模型一开始是用的avg_pool，如下：

self.avg_pool = qnn.QuantAvgPool2d(kernel_size=(1, 13), stride=1)

在训练、推断的时候都没问题，但是在下一步部署到FINN框架的时候会报错。这个最后没有解决，因此换成了max_pool，好处是不需要在init函数里定义这个层，只用在forward函数里按照torch最开始的方式写就行了，如下：

out = F.max_pool2d(out, kernel_size=(1, 13), stride=1)

猜想是max_pool 之前和之后的特征图不需要重新计算scale，所以可以用上面这种方式。

6. depth-wise & point-wise convlution

在brevitas框架下，亲测深度可分离卷积不大好用。在这个框架下这种卷积方式导致模型train不起来（亲测，换成正常的卷积就好用）。而且如果你后续要适配FINN，最好还是先尝试正常的卷积，因为深度可分离的卷积在FINN框架的默认配置下也会transform失败。。需要自己更改配置。

7. Residual Block

如果的模型有类似resnet的skip connection，那默认的buidl_dataflow是不生效的。你需要自己参考resnet50的dataflow. 见这个issue。

unfortunately the default build_dataflow steps will not work for a
model with residual connections. Maybe you can take inspiration from
the ResNet-50 finn-example on how to adjust the streamlining and
convert_to_hls steps accordingly. You can find the code here:

https://github.com/Xilinx/finn-examples/blob/main/build/resnet50/build.py
https://github.com/Xilinx/finn-examples/blob/main/build/resnet50/custom_steps.py

未完待续

就像开头提到的，我们最后有一个问题没有解决。在导出IP后，最后一步上板验证，发现每次计算出的结果不一样。这个问题官方也还没有解决。。所以后面我们用brevitas的训练，自己写了一个网络的推断过程。见下一篇。