卷积层和归一化层融合以及ONNX的使用

      之前就已经提到过，在神经网络种涉及到归一化的操作中就要特别小心。一方面是训练和推理阶段要有明确的标识来进行区分。否则，就会导致训练和推理的结果相差很大。另一方面就是归一化的方法很多，大致可分为：batchNorm、layerNorm、instanceNorm、GroupNorm、SwitchableNorm，他们直接的具体区别可以参考这里。

     之前常用的就是BN了，而且大都是选择默认的平台接口。偶然间的一次接触instanceNorm，让我重点关注到了其中的连个参数affine，track_running_stats，这两个参数默认都是True的状态。其具体用法可以简洁的介绍一下。我们知道归一化的操作一般是从bn发展开来的。每一层归一化后涉及到四个参数：weight（，阿尔法），bias（, 贝塔），running_mean、running_var。他们四个具体意义不再介绍。

其中affine控制着，。track_running_stats控制着running_mean、running_var。其实之前提到的训练和推理阶段的结果不一致的问题就是由于running_mean、running_var这两个参数造成的。因为在训练阶段，这个是根据归一化的方法来求特征层的均值和方差，并且根据一定的规则进行保存，最终目的是整个训练集都可以适应保存的这个值，从而使效果更好。如果在推理阶段没有明确的标识，从而导致识别成了训练阶段，那么这里的均值和方差就是根据输入的图片生成的特征层的值得到的。如果正确的标记成推理阶段，这里用到的方差是训练阶段保存的。

由上面提到的affine，track_running_stats可以看出，归一化可以选择是否使用缩放系数（ ， ）和均值方差。

那么，为什么会涉及到了卷积层和BN层的融合呢，也可以称之为merge。主要是使用归一化一般都是在卷积之后，而且他们之间是解偶的，如果合并成一个操作，一方面减少了变量数量，使模型小一些；另一方面，减少了计算量。其实一般的部署方案中涉及到的加速，这也是常用方法。关于bn的融合方法，已经有很多博客或者github给出结果和代码。例如这里、这里。

之所以这里再次重复介绍，主要是在pytorch-onnx模型转换过程种，使用instanceNorm并且affine=false，track_running_stats=true时，均值和方差无法正确的加载进去。转换后的结果与torch在training状态下的推理结果一样，通过这个错误就应当想到running_mean、running_var没有正确加载，然后在转换的过程种已经百分之百确定执行了model.eval()的操作。通过Netron查看转换成onnx的模型结构，也显示出了op名称是instanceNorm，weight（，阿尔法）与 bias（, 贝塔）其值均为1和0,是符合affine=false的状态。但是无法显示running_mean、running_var。因此把instanceNorm和卷积进行融合去掉归一化层，可以解决此问题。根据设置的特性，只需将bn融合的代码中weight（，阿尔法）与 bias（, 贝塔）改成1和0即可。

融合方法经历如下：首先打印出torch保存的键值对，由于affine是false的状态，卷积过后直接就是running_mean、running_var。没有bn.weight的相关key。总结出来规律提取涉及到的卷积权重和归一化层的权重即可。如果想要merge其他归一化的方法。根据其原理合状态进行修改代码即可

import torch
from collections import OrderedDict

def merge_bn(weights, bias, bn_mean, bn_var):
    weights_data = weights.data
    bias_data = bias.data
    bn_mean_data = bn_mean.data
    bn_var_data = bn_var.data
    weights_new = weights_data / torch.sqrt(bn_var_data.view(weights_data.size()[0], 1, 1, 1) + 1e-5)
    bias_new = (bias_data - bn_mean_data) / torch.sqrt(bn_var_data + 1e-5)
    return weights_new, bias_new

def merge_model():

    pt_model_path = 'old.pth'
    save_path = 'new.pth'
    trained_weights = torch.load(pt_model_path, map_location=torch.device('cuda'))
    load_net_clean = OrderedDict()

    global weights, bias
    global bn_mean, bn_var
    global weights_key, bias_key

    for name, params in trained_weights.items():

        if '.0.weight' in name:
            weights_key = name
            weights = params
        elif '.0.bias' in name:
            bias_key = name
            bias = params
        elif '.1.running_mean' in name:
            bn_mean = params
        elif '.1.running_var' in name:
            bn_var = params
            weights_new, bias_new = merge_bn(weights, bias, bn_mean, bn_var)
            load_net_clean[weights_key] = weights_new
            load_net_clean[bias_key] = bias_new
        elif '.1.num_batches_tracked' in name:
            continue
        else:
            load_net_clean[name] = params
    torch.save(load_net_clean, save_path)

if __name__ == '__main__':
    merge_model()

Torch-onnx

之所以使用模型转换，这是因为在推理阶段经常有专门对这一方面进行优化的平台。这使得在专门的平台进行部署，可以使推理速度有明显的提升。tensorrt就是一个很好的推理阶段部署平台。然后由于平台众多（pytorch，tensorflow，keras等），同一种方法内部实现并不一定相同，例如torch的特征层的的分布是[B，C，W，H]，但是在tf种C通道在最后一位。因此，不通平台按照某个标准转换到同一个平台下，这样就可以统一部署。从其他平台到tensorrt，一般都是转换成onnx。pytorch已经有专门的转换接口

torch.onnx.export(model.module, input_tensor, onnx_model_path, verbose=True, input_names=input_names,output_names=output_names, dynamic_axes=dynamic_axes)

方法的第一个参数是torch的模型，input_tensor为输入的tensor，即使是模型转换，也需要提供一个tensor，并且还有静态动态的输入尺寸的区别，这个是由dynamic_axes参数决定，如果是静态尺寸，在推理阶段的图片输入尺寸是不可以改变的。onnx_model_path是转换后模型的路径。verbose的作用是在转换过程种是否打印转换的op操作。

import torch
import torch.onnx
import onnx
from torch.autograd import Variable
import numpy as np
from collections import OrderedDict

pt_model_path = 'torch.pth'
onnx_model_path = 'onnx.onnx'

model = get_torch_model()

load_net = torch.load(pt_model_path, map_location=torch.device('cuda'))
model.load_state_dict(load_net)
model.load_state_dict(torch.load(pt_model_path))
model = model.eval()

def static_iputsize_trans():
    with torch.no_grad():
        img_input = np.ones((1080, 1920, 3))
        img_tensor = torch.from_numpy(np.transpose(img_input / 255.0, (2, 0, 1)).astype('float32'))
        img_tensor = Variable(img_tensor.unsqueeze(0).cuda())
        input_names = ['input']
        output_names = ['output']
        model.train(False)
        torch.onnx.export(model.module, img_tensor, onnx_model_path, verbose=True, input_names=input_names,
                          output_names=output_names)

        onnx.checker.check_model(onnx_model_path)
        print("==> Passed")

def dynamic_inputsize_trans():
    with torch.no_grad():
        input_tensor = torch.randn(1, 3, 512, 512).cuda()
        input_names = ['input']
        output_names = ['output']
        dynamic_axes = {'input': {0: 'batch', 2: 'height', 3: 'width'}, 'output': {0: 'batch', 2: 'height', 3: 'width'}}
        torch.onnx.export(model.module, input_tensor, onnx_model_path, verbose=True, input_names=input_names,
                          output_names=output_names, dynamic_axes=dynamic_axes)
        onnx.checker.check_model(onnx_model_path)
        print("==> Passed")

if __name__ == '__main__':
    dynamic_inputsize_trans()

其中创建torch模型方法get_torch_model()的返回过程一般都加入到

nn.DataParallel(model)

因此在torch.onnx.export种的第一个参数才加入了module，即model.module。

tf-onnx

tensorflow好像没有直接提供类似torch的接口。但是可以通过pip安装tf2onnx这个工具包。安装好了就可以直接使用python调用了

python -m tf2onnx.convert --input input.pb  --inputs input_1:0,input_2:0,input_3:0 --outputs out_1:0,out_2:0,out_3:0 --output output.onnx --opset 13 --verbose 

 --opeset 13   添加这个后支持的OP就更加多一些。

但是tensorflow训练出来的通道一般在-1的位置，和torch不一样，转换的话可以这样

<没有设置>
python -m tf2onnx.convert --input input.pb  --inputs input_1:0,input_2:0,input_3:0 --outputs out_1:0,out_2:0,out_3:0 --output output.onnx --opset 13 --verbose --inputs-as-nchw  input_1:0,input_2:0,input_3:0
 

onnx推理

验证onnx模型是否正确转换

import onnxruntime as ort
import cv2
import numpy as np
def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    ort_session = ort.InferenceSession(onnx_model_path)
    img = cv2.imread(img_path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
    img = np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))
    img = np.expand_dims(img, 0)
    output = ort_session.run(None, {'input': img.astype(np.float32)})
    output = output[0]
    output = np.squeeze(output, 0)
    output = np.clip(output, 0, 1)
    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
    output = np.clip((output * 255.0), 0, 255).astype(np.uint8)
    cv2.imwrite(save_path, output)

pycuda

接下来可以安装tensorrt，通过加载onnx来验证tensorrt平台的性能。具体的推理方法不再详细介绍，主要是在安装pycuda遇到的问题给出解决方法。

使用pip安装pycuda报错 src/cpp/cuda.hpp:14:18: fatal error: cuda.h: No such file or directory

这之前或许使用了sudo pip3 install pycuda。如果不使用sudo，无法获取权限。但是使用了，cuda内部也是无法拥有root权限。解决办法：

sudo su export PATH=/usr/local/cuda-10.0/bin:/usr/local/cuda/bin:$PATH

然后再次使用
pip3 install pycuda 命令