成功解决Pytorch模型转trt模型中得BatchNorm问题

问题背景 如果你想把你的模型投入到应用中或者是想提升模型的运行速度，除了对网络进行压缩、蒸馏外，最好的方法就是将模型转成tensor模型，使用tensorrt实现对网络的加速。但是当该模型的功能是图像增强或者是图像生成，并且模型中运用了大量的batchnorm2d函数，运用网上现成的方法会发现模型转成onnx以及trt后，模型的处理效果大幅下降，想解决此问题就可以详细往下看了： 我们的方法顺序是：pytorch模型先转成onnx模型，接着将onnx模型转成trt模型 一、pytorch to onnx 核心代码： import torch from torchvision.utils import save_image import os from nets.tiny_unet_2_channelxian..