pytorch dropout_pytorch转onnx中batchnorm的坑

这个帖子最初只是用以安慰自己这几天白忙活的心情，没想到居然能够引来各位大佬在评论区答疑解惑，受益良多，也正符合了我将自己这些不成熟的经历写出来的初衷，哪怕只有一个人在遇到类似问题的时候能搜索到这篇文章，给之以帮助，就很满足了。也希望各位读者若是看完文章的同时能更仔细的浏览一下大家的评论，相信会有更深的理解。

---

怀着难以压抑的心情记录下这段坎坷经历。

前一晚还沉浸在成功将pytorch模型转成onnx并部署在tensorrt上，实现了肉眼可见的速度提升，而且还支持动态尺寸输入，踏踏实实的睡了一觉，醒来之后在大规模图像上测试战果，结果心情直接跌落谷底，tensorrt预测输出和pytorch结果有明显出入，只有少数情况下一致，下面简单记录一下debug的过程，也希望能帮助到后来者。

整个过程经历了两次格式的转换，正经的严谨的思路是每次转换完之后都应该去确认一下转换前后结果是否一致，奈何我偏偏挑了一条少数情况下一致的数据，结果最后才发现问题。

最开始的时候以为是onnx转tensorrt时报的warning导致出错，因为onnx的格式是INT64，而tensorrt只支持INT32的格式，是不是这个过程中有些精度损失，可是这损失也太大了，幸好我尝试了用onnxruntime载入onnx文件直接推理，结果却和tensorrt完全一样，好的时候大家都好，不好的时候谁也没落下，那说明不是tensorrt的问题。

onnx转tensorrt没有问题的话，那么只能是pytorch转onnx的问题，实际情况也确实是这样，pytoch明明预测的挺好，咋到了onnx就不行呢，于是我对比了一下pytorch模型和onnx模型参数的区别，总不会参数一样输出却不相同吧。恭喜你，还真是这样！pytoch打印出来的model参数和netron里onnx的参数一模一样，幸好我又尝试了一下喂入图像，挑选结果有差别的图像分别在pytorch和onnx上推理，结果推理完参数果然有变化，onnx模型推理前后参数保持一致，而pytorch模型推理前后batchnorm层变得不一样了。

倒推的过程大概就是这样，废话太多了，我就直接抛结论了。

pytorch转onnx前后推理结果有明显差别的根源是model.train()和model.eval()的差别，之前其实发现了eval模式下模型的表现不好，不过却没有深究，结果每一道坎你以为跳过去的，但总会在后面的路悄悄坑你一把。这两种模式的主要区别是batchnorm和dropout的区别，dropout不是今天的主角，主要讲讲batchnorm的区别。train模式和eval模式下，batchnorm的区别有很多，可以参见各位大佬的科普，但是造成转onnx之后结果差异如此之大的原因是eval模式下，batchnorm的mean和var是固定的，固定的啊同志们，不管你的输入是啥，我都是固定的，最开始我单纯的以为batchnorm、layernorm等等norm都是做了一个类似减去期望处以标准差的归一化操作（实际上主要思想确实是），但是不管啥输入，你的统计结果都一样是什么鬼，虽然这个统计结果是基于训练集整体数据的统计吧，但是万一我推理的时候数据有些不一样呢？？？在eval模式下，咱们依然可以改成参数随着输入更新，用下面几行代码就可以，但是导出onnx的时候它可不管你是train模式还是eval模式，直接以eval模式的设置导出啊，不管你怎么改，我就是batchnorm不更新啊，我就是认为你的测试数据和训练数据严格同分布啊！

def unfreeze_bn(model):
    for m in model.modules():
            if isinstance(m,nn.BatchNorm2d):
                m.train()

虽然能够理解onnx为了一些速度和一致性的考虑下，固定batchnorm的参数，但是还是有些惋惜浪费了这么长时间去debug，要怪只怪当初怎么不多选点数据在pytorch转onnx的时候就把这个问题找出来，白瞎后面投入不少精力的tensorrt动态部署，最终也没能成功部署，只能看看在数据读取上加快速度了。