071、超分模型部署ONNX 导出与 TensorRT 加速实战上周帮同事调一个EDSR模型的部署问题模型在PyTorch里跑得挺欢一上TensorRT就崩报错信息看得人头皮发麻——“Unsupported layer type”。翻日志发现是PixelShuffle那个上采样操作在TensorRT里没被正确解析。这种坑我踩了不下十次今天干脆把超分模型从ONNX导出到TensorRT加速的完整流程和那些容易翻车的地方都记下来。从PyTorch到ONNX第一步就卡住很多人觉得ONNX导出就是torch.onnx.export一行代码的事天真了。超分模型里那些自定义操作、动态尺寸、上采样层每一个都是潜在的雷。先看一个最基础的EDSR导出代码注意我踩过的坑importtorchimporttorch.onnx# 假设你有一个训练好的EDSR模型modelEDSR().eval()# 别这样写固定尺寸导出会让你后面哭dummy_inputtorch.randn(1,3,224,224)# 固定尺寸大忌# 正确的做法用动态轴dummy_inputtorch.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_version11,# 这里踩过坑opset太低不支持某些操作do_constant_foldingTrue,verboseFalse)dynamic_axes这个参数是超分模型的生命线。超分任务输入尺寸千变万化固定尺寸导出意味着你只能处理一种分辨率那还叫什么超分但注意动态轴会带来额外的性能开销后面TensorRT优化的时候要权衡。opset版本我建议至少用11低于这个版本像torch.nn.functional.interpolate这种操作可能导出失败。如果你用了nn.PixelShuffleopset 11以上才能保证正确导出。那些让你抓狂的ONNX导出错误错误1RuntimeError: 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_foldingTrue搞的鬼。这个参数会把一些常量计算折叠到权重里但超分模型里有些操作比如归一化层的均值和方差被折叠后导出的ONNX在推理时行为就变了。我的经验先设成False导出验证结果正确后再开True做优化。ONNX模型验证别跳过这一步导出完ONNX别急着转TensorRT。先用ONNX Runtime跑一遍确认数值精度没问题importonnxruntimeasortimportnumpyasnp# 创建ONNX Runtime sessionsessort.InferenceSession(edsr.onnx)# 准备输入数据用和PyTorch一样的tensorinput_datanp.random.randn(1,3,64,64).astype(np.float32)# 推理outputssess.run(None,{input:input_data})# 和PyTorch结果对比withtorch.no_grad():pt_outputmodel(torch.from_numpy(input_data)).numpy()# 计算相对误差rel_errornp.abs(outputs[0]-pt_output)/(np.abs(pt_output)1e-8)print(fMax relative error:{rel_error.max():.6f})相对误差超过1e-5就要警惕了。我遇到过一种情况ONNX和PyTorch结果差了0.1%查了两天发现是torch.onnx.export里的training参数没设成False导致BatchNorm层行为不一致。TensorRT加速从ONNX到EngineONNX验证通过后就该上TensorRT了。这里我推荐用trtexec命令行工具先快速验证再写代码集成。# 先用trtexec看看能不能转成功trtexec--onnxedsr.onnx--saveEngineedsr_fp32.engine--workspace4096# 如果显存够试试FP16trtexec--onnxedsr.onnx--saveEngineedsr_fp16.engine--fp16--workspace4096--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)]:dummytorch.randn(1,3,h,w)torch.onnx.export(model,dummy,fedsr_{h}x{w}.onnx,input_names[input],output_names[output],opset_version11)然后每个尺寸生成一个engine推理时根据输入尺寸选择对应的engine。缺点很明显内存占用大而且不能处理任意尺寸。但胜在稳定工业部署里很多人这么干。实战用Python API构建TensorRT推理假设你已经成功生成了engine文件下面是推理代码注意那些容易翻车的地方importtensorrtastrtimportpycuda.driverascudaimportpycuda.autoinitimportnumpyasnpclassTRTInference:def__init__(self,engine_path):# 别这样写每次推理都创建runtime和context# 应该复用self.loggertrt.Logger(trt.Logger.WARNING)self.runtimetrt.Runtime(self.logger)# 读取engine文件withopen(engine_path,rb)asf:self.engineself.runtime.deserialize_cuda_engine(f.read())self.contextself.engine.create_execution_context()# 获取输入输出信息self.input_idxself.engine.get_binding_index(input)self.output_idxself.engine.get_binding_index(output)# 分配显存这里踩过坑要按最大尺寸分配self.max_input_size1*3*256*256*4# 假设最大输入256x256self.max_output_size1*3*1024*1024*4# 4倍超分self.d_inputcuda.mem_alloc(self.max_input_size)self.d_outputcuda.mem_alloc(self.max_output_size)# 固定尺寸的输入输出host内存self.h_inputnp.empty((1,3,64,64),dtypenp.float32)self.h_outputnp.empty((1,3,256,256),dtypenp.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_shapeself.context.get_binding_shape(self.output_idx)self.h_outputnp.empty(output_shape,dtypenp.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_outputset_binding_shape这个调用是动态尺寸engine的关键。很多人忘了这一步结果推理出来的图像尺寸不对或者直接报错。另外注意execute_v2的bindings参数传的是显存地址的int值不是numpy数组这个坑我见过新手栽了无数次。性能调优从能跑到跑得快engine能跑了但速度不满意几个调优方向1. 选择合适的数据精度FP16理论上能带来2倍加速但超分模型对精度敏感。我测试过EDSRFP16下PSNR掉了0.3dB肉眼看不出来但论文里写的就是0.3dB的差距。如果对精度要求苛刻用INT8量化但需要校准数据集而且超分任务的量化难度比分类大得多。2. 调整workspace大小trtexec或build_engine时的workspace参数直接影响TensorRT能尝试的tactic数量。workspace越大越可能找到更快的kernel。我一般设到显存上限的80%。3. 使用多流推理视频超分场景下可以用多个CUDA stream并行处理不同帧stream1cuda.Stream()stream2cuda.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上。那些年我踩过的坑最后列几个血泪教训希望能帮你少走弯路坑1BatchNorm和InstanceNorm在推理时的行为超分模型训练时用BatchNorm但推理时BN层的行为会变。ONNX导出时一定要确保模型在eval模式否则导出的ONNX里BN层会保留训练时的统计量推理结果直接崩。坑2动态尺寸的显存分配动态尺寸意味着你无法预知输入大小。我见过有人每次推理都重新分配显存结果速度比PyTorch还慢。正确做法按最大可能尺寸一次性分配或者用显存池。坑3TensorRT版本兼容性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_modeTrue的池化层那输出尺寸的计算逻辑会变得很复杂TensorRT可能算错。我的建议所有池化层都用ceil_modeFalse或者干脆用卷积替代。个人经验总结超分模型部署这件事说难不难说简单也不简单。核心就三点ONNX导出要稳、TensorRT转换要准、推理代码要快。每一步都有坑但踩过一次就记住了。如果你刚开始做这件事我的建议是先别追求动态尺寸固定几个常用分辨率比如720p、1080p做多engine部署。等流程跑通了再考虑动态尺寸的优化。毕竟在工业场景里稳定比灵活更重要。另外别迷信TensorRT的自动优化。有时候手动写一个简单的CUDA kernel替换某些操作效果反而更好。比如PixelShuffleTensorRT的实现在某些GPU上有bug我自己写了一个kernel速度提升了30%。最后调试工具要会用。trtexec的--dumpProfile参数能告诉你每个层的耗时onnxruntime的InferenceSession可以设置enable_profilingTrue。性能瓶颈在哪一目了然。好了就写这么多。下次遇到超分模型部署的问题欢迎来评论区交流。
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模型modelEDSR().eval()# 别这样写固定尺寸导出会让你后面哭dummy_inputtorch.randn(1,3,224,224)# 固定尺寸大忌# 正确的做法用动态轴dummy_inputtorch.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_version11,# 这里踩过坑opset太低不支持某些操作do_constant_foldingTrue,verboseFalse)dynamic_axes这个参数是超分模型的生命线。超分任务输入尺寸千变万化固定尺寸导出意味着你只能处理一种分辨率那还叫什么超分但注意动态轴会带来额外的性能开销后面TensorRT优化的时候要权衡。opset版本我建议至少用11低于这个版本像torch.nn.functional.interpolate这种操作可能导出失败。如果你用了nn.PixelShuffleopset 11以上才能保证正确导出。那些让你抓狂的ONNX导出错误错误1RuntimeError: 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_foldingTrue搞的鬼。这个参数会把一些常量计算折叠到权重里但超分模型里有些操作比如归一化层的均值和方差被折叠后导出的ONNX在推理时行为就变了。我的经验先设成False导出验证结果正确后再开True做优化。ONNX模型验证别跳过这一步导出完ONNX别急着转TensorRT。先用ONNX Runtime跑一遍确认数值精度没问题importonnxruntimeasortimportnumpyasnp# 创建ONNX Runtime sessionsessort.InferenceSession(edsr.onnx)# 准备输入数据用和PyTorch一样的tensorinput_datanp.random.randn(1,3,64,64).astype(np.float32)# 推理outputssess.run(None,{input:input_data})# 和PyTorch结果对比withtorch.no_grad():pt_outputmodel(torch.from_numpy(input_data)).numpy()# 计算相对误差rel_errornp.abs(outputs[0]-pt_output)/(np.abs(pt_output)1e-8)print(fMax relative error:{rel_error.max():.6f})相对误差超过1e-5就要警惕了。我遇到过一种情况ONNX和PyTorch结果差了0.1%查了两天发现是torch.onnx.export里的training参数没设成False导致BatchNorm层行为不一致。TensorRT加速从ONNX到EngineONNX验证通过后就该上TensorRT了。这里我推荐用trtexec命令行工具先快速验证再写代码集成。# 先用trtexec看看能不能转成功trtexec--onnxedsr.onnx--saveEngineedsr_fp32.engine--workspace4096# 如果显存够试试FP16trtexec--onnxedsr.onnx--saveEngineedsr_fp16.engine--fp16--workspace4096--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)]:dummytorch.randn(1,3,h,w)torch.onnx.export(model,dummy,fedsr_{h}x{w}.onnx,input_names[input],output_names[output],opset_version11)然后每个尺寸生成一个engine推理时根据输入尺寸选择对应的engine。缺点很明显内存占用大而且不能处理任意尺寸。但胜在稳定工业部署里很多人这么干。实战用Python API构建TensorRT推理假设你已经成功生成了engine文件下面是推理代码注意那些容易翻车的地方importtensorrtastrtimportpycuda.driverascudaimportpycuda.autoinitimportnumpyasnpclassTRTInference:def__init__(self,engine_path):# 别这样写每次推理都创建runtime和context# 应该复用self.loggertrt.Logger(trt.Logger.WARNING)self.runtimetrt.Runtime(self.logger)# 读取engine文件withopen(engine_path,rb)asf:self.engineself.runtime.deserialize_cuda_engine(f.read())self.contextself.engine.create_execution_context()# 获取输入输出信息self.input_idxself.engine.get_binding_index(input)self.output_idxself.engine.get_binding_index(output)# 分配显存这里踩过坑要按最大尺寸分配self.max_input_size1*3*256*256*4# 假设最大输入256x256self.max_output_size1*3*1024*1024*4# 4倍超分self.d_inputcuda.mem_alloc(self.max_input_size)self.d_outputcuda.mem_alloc(self.max_output_size)# 固定尺寸的输入输出host内存self.h_inputnp.empty((1,3,64,64),dtypenp.float32)self.h_outputnp.empty((1,3,256,256),dtypenp.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_shapeself.context.get_binding_shape(self.output_idx)self.h_outputnp.empty(output_shape,dtypenp.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_outputset_binding_shape这个调用是动态尺寸engine的关键。很多人忘了这一步结果推理出来的图像尺寸不对或者直接报错。另外注意execute_v2的bindings参数传的是显存地址的int值不是numpy数组这个坑我见过新手栽了无数次。性能调优从能跑到跑得快engine能跑了但速度不满意几个调优方向1. 选择合适的数据精度FP16理论上能带来2倍加速但超分模型对精度敏感。我测试过EDSRFP16下PSNR掉了0.3dB肉眼看不出来但论文里写的就是0.3dB的差距。如果对精度要求苛刻用INT8量化但需要校准数据集而且超分任务的量化难度比分类大得多。2. 调整workspace大小trtexec或build_engine时的workspace参数直接影响TensorRT能尝试的tactic数量。workspace越大越可能找到更快的kernel。我一般设到显存上限的80%。3. 使用多流推理视频超分场景下可以用多个CUDA stream并行处理不同帧stream1cuda.Stream()stream2cuda.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上。那些年我踩过的坑最后列几个血泪教训希望能帮你少走弯路坑1BatchNorm和InstanceNorm在推理时的行为超分模型训练时用BatchNorm但推理时BN层的行为会变。ONNX导出时一定要确保模型在eval模式否则导出的ONNX里BN层会保留训练时的统计量推理结果直接崩。坑2动态尺寸的显存分配动态尺寸意味着你无法预知输入大小。我见过有人每次推理都重新分配显存结果速度比PyTorch还慢。正确做法按最大可能尺寸一次性分配或者用显存池。坑3TensorRT版本兼容性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_modeTrue的池化层那输出尺寸的计算逻辑会变得很复杂TensorRT可能算错。我的建议所有池化层都用ceil_modeFalse或者干脆用卷积替代。个人经验总结超分模型部署这件事说难不难说简单也不简单。核心就三点ONNX导出要稳、TensorRT转换要准、推理代码要快。每一步都有坑但踩过一次就记住了。如果你刚开始做这件事我的建议是先别追求动态尺寸固定几个常用分辨率比如720p、1080p做多engine部署。等流程跑通了再考虑动态尺寸的优化。毕竟在工业场景里稳定比灵活更重要。另外别迷信TensorRT的自动优化。有时候手动写一个简单的CUDA kernel替换某些操作效果反而更好。比如PixelShuffleTensorRT的实现在某些GPU上有bug我自己写了一个kernel速度提升了30%。最后调试工具要会用。trtexec的--dumpProfile参数能告诉你每个层的耗时onnxruntime的InferenceSession可以设置enable_profilingTrue。性能瓶颈在哪一目了然。好了就写这么多。下次遇到超分模型部署的问题欢迎来评论区交流。
 获取与实战)
双梁正问题 【 torch 实战】（Python代码实现）)
