071、超分模型部署：ONNX 导出与 TensorRT 加速实战-拓冰建站

071、超分模型部署：ONNX 导出与 TensorRT 加速实战

上周帮同事调一个EDSR模型的部署问题，模型在PyTorch里跑得挺欢，一上TensorRT就崩，报错信息看得人头皮发麻——“Unsupported layer type”。翻日志发现是PixelShuffle那个上采样操作在TensorRT里没被正确解析。这种坑我踩了不下十次，今天干脆把超分模型从ONNX导出到TensorRT加速的完整流程和那些容易翻车的地方都记下来。

从PyTorch到ONNX，第一步就卡住

很多人觉得ONNX导出就是torch.onnx.export一行代码的事，天真了。超分模型里那些自定义操作、动态尺寸、上采样层，每一个都是潜在的雷。

先看一个最基础的EDSR导出代码，注意我踩过的坑：

importtorchimporttorch.onnx# 假设你有一个训练好的EDSR模型model=EDSR().eval()# 别这样写！固定尺寸导出会让你后面哭dummy_input=torch.randn(1,3,224,224)# 固定尺寸，大忌# 正确的做法：用动态轴dummy_input=torch.randn(1,3,64,64)# 小尺寸就行，关键是动态轴设置torch.onnx.export(model,dummy_input,"edsr.onnx",input_names=['input'],output_names=['output'],dynamic_axes={'input':{2:'height',3:'width'},'output':{2:'height',3:'width'}},opset_version=11,# 这里踩过坑，opset太低不支持某些操作do_constant_folding=True,verbose=False)

dynamic_axes这个参数是超分模型的生命线。超分任务输入尺寸千变万化，固定尺寸导出意味着你只能处理一种分辨率，那还叫什么超分？但注意，动态轴会带来额外的性能开销，后面TensorRT优化的时候要权衡。

opset版本我建议至少用11，低于这个版本，像torch.nn.functional.interpolate这种操作可能导出失败。如果你用了nn.PixelShuffle，opset 11以上才能保证正确导出。

那些让你抓狂的ONNX导出错误

错误1：RuntimeError: Exporting the operator xxx to ONNX opset version 11 is not supported

八成是你用了太新的PyTorch操作。超分模型里常见的F.interpolate、PixelShuffle、DepthToSpace这些，在opset 11里都有对应映射。但如果你用了torch.nn.functional.grid_sample做变形校正，那得小心，这个操作在opset 16以后才稳定。

解决方案：要么升级opset，要么用torch.onnx.register_custom_op_symbolic手写映射。我一般选前者，后者写起来像在给编译器写遗嘱。

错误2：导出成功但推理结果不对

这个最阴间。模型能跑，但输出全是噪声。大概率是do_constant_folding=True搞的鬼。这个参数会把一些常量计算折叠到权重里，但超分模型里有些操作（比如归一化层的均值和方差）被折叠后，导出的ONNX在推理时行为就变了。

我的经验：先设成False导出，验证结果正确后再开True做优化。

ONNX模型验证，别跳过这一步

导出完ONNX，别急着转TensorRT。先用ONNX Runtime跑一遍，确认数值精度没问题：

importonnxruntimeasortimportnumpyasnp# 创建ONNX Runtime sessionsess=ort.InferenceSession("edsr.onnx")# 准备输入数据，用和PyTorch一样的tensorinput_data=np.random.randn(1,3,64,64).astype(np.float32)# 推理outputs=sess.run(None,{'input':input_data})# 和PyTorch结果对比withtorch.no_grad():pt_output=model(torch.from_numpy(input_data)).numpy()# 计算相对误差rel_error=np.abs(outputs[0]-pt_output)/(np.abs(pt_output)+1e-8)print(f"Max relative error:{rel_error.max():.6f}")

相对误差超过1e-5就要警惕了。我遇到过一种情况：ONNX和PyTorch结果差了0.1%，查了两天发现是torch.onnx.export里的training参数没设成False，导致BatchNorm层行为不一致。

TensorRT加速：从ONNX到Engine

ONNX验证通过后，就该上TensorRT了。这里我推荐用trtexec命令行工具先快速验证，再写代码集成。

# 先用trtexec看看能不能转成功trtexec--onnx=edsr.onnx--saveEngine=edsr_fp32.engine--workspace=4096# 如果显存够，试试FP16trtexec--onnx=edsr.onnx--saveEngine=edsr_fp16.engine--fp16--workspace=4096

--workspace参数设大点，超分模型参数量大，workspace太小会报"out of memory"。我一般设4GB起步，8GB更稳。

如果trtexec能成功生成engine文件，恭喜你，最难的坎过去了。但现实往往是：

[ERROR] Some tactics do not have sufficient workspace memory to run. [ERROR] Cannot find a valid tactic for layer: Resize_123

又是上采样层的问题。TensorRT对Resize操作的支持有限，尤其是动态尺寸下的Resize。解决方案有两个：

方案一：用TensorRT的IResizeLayer替换ONNX里的Resize

这个方案需要你写TensorRT的C++或Python API，手动构建网络。我一般只在模型结构简单的时候用，EDSR这种还好，RCAN那种残差块堆叠的，手写网络能写到怀疑人生。

方案二：在PyTorch里用固定尺寸导出，然后做多尺寸engine

# 导出多个固定尺寸的ONNXforh,win[(64,64),(128,128),(256,256)]:dummy=torch.randn(1,3,h,w)torch.onnx.export(model,dummy,f"edsr_{h}x{w}.onnx",input_names=['input'],output_names=['output'],opset_version=11)

然后每个尺寸生成一个engine，推理时根据输入尺寸选择对应的engine。缺点很明显：内存占用大，而且不能处理任意尺寸。但胜在稳定，工业部署里很多人这么干。

实战：用Python API构建TensorRT推理

假设你已经成功生成了engine文件，下面是推理代码，注意那些容易翻车的地方：

importtensorrtastrtimportpycuda.driverascudaimportpycuda.autoinitimportnumpyasnpclassTRTInference:def__init__(self,engine_path):# 别这样写：每次推理都创建runtime和context# 应该复用self.logger=trt.Logger(trt.Logger.WARNING)self.runtime=trt.Runtime(self.logger)# 读取engine文件withopen(engine_path,'rb')asf:self.engine=self.runtime.deserialize_cuda_engine(f.read())self.context=self.engine.create_execution_context()# 获取输入输出信息self.input_idx=self.engine.get_binding_index('input')self.output_idx=self.engine.get_binding_index('output')# 分配显存，这里踩过坑：要按最大尺寸分配self.max_input_size=1*3*256*256*4# 假设最大输入256x256self.max_output_size=1*3*1024*1024*4# 4倍超分self.d_input=cuda.mem_alloc(self.max_input_size)self.d_output=cuda.mem_alloc(self.max_output_size)# 固定尺寸的输入输出host内存self.h_input=np.empty((1,3,64,64),dtype=np.float32)self.h_output=np.empty((1,3,256,256),dtype=np.float32)definfer(self,input_np):# 检查输入尺寸，动态engine需要设置输入形状ifself.engine.has_implicit_batch_dimension:# 固定batch的enginepasselse:# 动态shape的engine，必须设置self.context.set_binding_shape(self.input_idx,input_np.shape)# 获取实际输出形状output_shape=self.context.get_binding_shape(self.output_idx)self.h_output=np.empty(output_shape,dtype=np.float32)# 拷贝数据到显存cuda.memcpy_htod(self.d_input,input_np.astype(np.float32))# 执行推理self.context.execute_v2(bindings=[int(self.d_input),int(self.d_output)])# 拷贝结果回hostcuda.memcpy_dtoh(self.h_output,self.d_output)returnself.h_output

set_binding_shape这个调用是动态尺寸engine的关键。很多人忘了这一步，结果推理出来的图像尺寸不对，或者直接报错。另外注意，execute_v2的bindings参数传的是显存地址的int值，不是numpy数组，这个坑我见过新手栽了无数次。

性能调优：从能跑到跑得快

engine能跑了，但速度不满意？几个调优方向：

1. 选择合适的数据精度

FP16理论上能带来2倍加速，但超分模型对精度敏感。我测试过EDSR，FP16下PSNR掉了0.3dB，肉眼看不出来，但论文里写的就是0.3dB的差距。如果对精度要求苛刻，用INT8量化，但需要校准数据集，而且超分任务的量化难度比分类大得多。

2. 调整workspace大小

trtexec或build_engine时的workspace参数直接影响TensorRT能尝试的tactic数量。workspace越大，越可能找到更快的kernel。我一般设到显存上限的80%。

3. 使用多流推理

视频超分场景下，可以用多个CUDA stream并行处理不同帧：

stream1=cuda.Stream()stream2=cuda.Stream()# 交替提交推理任务context1.execute_async_v2(bindings1,stream1.handle)context2.execute_async_v2(bindings2,stream2.handle)

注意：每个stream需要独立的context，不能共享。

4. 避免频繁的CPU-GPU数据传输

这个是最容易被忽视的。如果你每帧都从CPU拷贝数据到GPU，再拷回去，那TensorRT的加速效果会被传输延迟吃掉。解决方案：用CUDA的pinned memory，或者干脆把预处理也搬到GPU上。

那些年我踩过的坑

最后列几个血泪教训，希望能帮你少走弯路：

坑1：BatchNorm和InstanceNorm在推理时的行为

超分模型训练时用BatchNorm，但推理时BN层的行为会变。ONNX导出时一定要确保模型在eval模式，否则导出的ONNX里BN层会保留训练时的统计量，推理结果直接崩。

坑2：动态尺寸的显存分配

动态尺寸意味着你无法预知输入大小。我见过有人每次推理都重新分配显存，结果速度比PyTorch还慢。正确做法：按最大可能尺寸一次性分配，或者用显存池。

坑3：TensorRT版本兼容性

TensorRT 7和8的API变化很大，尤其是动态shape的支持。如果你用TensorRT 7，动态尺寸支持有限，建议升级到8.x。另外，不同版本的TensorRT生成的engine文件不兼容，部署时一定要保证环境一致。

坑4：多GPU环境下的engine加载

engine文件是和GPU架构绑定的。你在A100上生成的engine，拿到V100上跑不了。如果有多卡环境，每张卡都要单独生成engine，或者用trtexec的--device参数指定。

坑5：超分模型的输出尺寸计算

超分模型的输出尺寸和输入尺寸、上采样倍数、padding方式都有关系。TensorRT里动态尺寸的engine，输出shape是推理时根据输入shape实时计算的。如果你在导出ONNX时用了ceil_mode=True的池化层，那输出尺寸的计算逻辑会变得很复杂，TensorRT可能算错。我的建议：所有池化层都用ceil_mode=False，或者干脆用卷积替代。