pytorch 语义分割模型先转ONNX再转NCNN

0 前言

最近项目需要用到物体分割思路进行开展，因为自己一直使用的是tensorflow1.x版本进行模型训练，故而分割模型训练之初便选用了tensorflow框架进行模型训练，整个训练过程是很顺利的，但项目落地过程中需要用到opencv或者ncnn前向推理框架进行实现，好家伙对于tf.pb模型那叫个调试的酸爽，对于最新版本的opencv(4.4.1)废了很大劲在模型加载过程中总算不报错了，但前向推理的结果却和tf环境下的前向推理结果不一致，那就继续改吧，改的是昏天黑地，死去活来，最终从纠结到放弃，因为opencv对于分割模型的支持实在有限，无奈只得放弃。反思有没有其他路径来解决这个问题呢，最后选择了这条路线进行实现pytorch->onnx->ncnn，看到此博客的你如果也在遭受分割模型转换的蹂躏，请静下心来按照我提供的这条路线一步一步走下去，你会走成功的。

1 pytorch

对于用tensorflow1.x已经习惯了的炼丹师，换成用pytorch不晓得有多爽，是真的很爽的那种，以至于目前抛弃了tf拥抱pytorch，首先呢就是将自己的分割网络改写成pytorch格式，至于怎么改自己百度学习去，这里需要提醒的是对于pytorch下的nn.MaxUnpool2d函数需要修改，可能你会纠结为啥要改呢？因为目前opencv和ncnn都不支持这个函数，所以不改不行，那改成啥呢？改成用nn.UpsamplingNearest2d，结果就是丢弃了最大值位置的索引，在模型精度上会有些影响，这里还要特别提醒对于pytorch网络结构中尽量别搞特殊的层，因为特殊层的风险就是opencv和ncnn不支持，比如在前期我用到了自定义的torch.zeros()这个函数为了拼接成需要的数据维度，然后ncnn就不支持它，无奈只能返回修改网络结构。

2 ONNX

pytorch训练的模型保存下来是以*.pth，要把它转换成ONNX，

import torch.onnx
torch.onnx.export(net, dummy_input, outputmodel,export_params=True,
                    input_names=input_names,output_names=output_names)

会生成*.onnx模型，然后用编译好的ncnn下的ncnn-master/build/tools/onnx下的onnx2ncnn.exe生成*.bin和*.param文件，结果报如下错误，

兄嘚，不要慌，因为*.onnx模型需要先进行onnx-simplifier一下，如果没有安装这个工具包就pip install 一哈子，simplifier命令行如下：

python -m onnxsim ./*.onnx ./*_sim.onnx

然后就会生成一个*_sim.onnx模型，拿这个模型去用编译好的ncnn工具生成*.bin和*.param文件就没问题了，模型生成完就用C语言写一个测试程序，然后选择一张图片看看它分别在pytorch环境下的结果与ncnn下的结果是否一致，如果一致就搞定了，如果不一致怎么办呢？也很好解决只不过得废些精力，需要根据*.param中提供的blob逐层输出前向结果，然后与pytorch下的逐层输出结果进行一一比对，定位到问题发生的点，然后再进行修改就阔以了，对于pytorch下如何逐层输出结果呢？给段代码仅供参考哈

x = self.transposed_conv(x, output_size=(360,480))
v = x.cpu().detach().numpy()
v1 = v.flatten()
fp = open('transposed_conv.txt', 'w')
fp.write(str(v.shape)+'\n')
for d in v1:
   fp.write(str(d)+'\n')
fp.close() 

可能有的娃子还在纠结如何逐层输出ncnn的结果，也给个参考代码哈

ncnn::Net plnet;
std::string param_files =  "../models/torch.param";
std::string bin_files = "../models/torch.bin";
plnet.load_param(param_files.data());
plnet.load_model(bin_files.data());

printf("load model success!!!\n");
cv::Mat srcImage = cv::imread("./*.jpg");	
ncnn::Mat ncnn_img = ncnn::Mat::from_pixels(srcImage.data, ncnn::Mat::PIXEL_BGR,         
                                                 srcImage.cols, srcImage.rows);
ncnn_img.substract_mean_normalize(mean_vals, norm_vals);
ncnn::Extractor ex = pupilnet.create_extractor();
//ex.set_num_threads(4);
//ex.set_light_mode(true);
ex.input("x", ncnn_img);
ncnn::Mat out;
ex.extract("y", out);

其中x就是模型的输入blob，y是模型的输出blob，如果不同层进行输出只修改y就可以了，这里要特别注意blob要写*.param中对应的键值，比如：Concat           Concat_2                 2 1 301 302 303 0=0 其中2表示有两个输入，1表示有一个输出，其中301 302表示输入的键值，303表示输出的键值，如果要输出Concat这个层的结果，就将y修改成303就阔以了

好了，博客就介绍到这了，祝愿你早日打通pytorch->onnx->ncnn分割模型转换，加油吧！炼丹师