047、移动端部署实战：ONNX导出、TensorRT加速与NCNN移植全攻略

047、移动端部署实战ONNX导出、TensorRT加速与NCNN移植全攻略昨天凌晨三点我盯着终端里那个“segmentation fault”的报错咖啡杯底已经结了厚厚一层咖啡渍。一个在PyTorch里跑得稳稳当当的EDSR模型导出ONNX后死活过不了TensorRT的校验最后在NCNN上直接崩了。这种场景做超分部署的兄弟应该都不陌生——训练时精度再高部署时踩的坑一个都不会少。今天这篇笔记我把从ONNX导出到NCNN移植的完整流程拆开揉碎重点讲那些文档里不会写的“暗坑”。你手头如果有现成的超分模型建议跟着走一遍别等到项目交付前三天才来翻这篇。一、ONNX导出别让torch.onnx.export坑了你很多人觉得ONNX导出就是一行代码的事结果导出的模型要么精度掉一截要么算子不支持。我踩过最深的坑是动态轴与静态轴的混淆。超分模型的输入通常是[1, 3, H, W]但H和W在推理时可能变化。如果你用torch.onnx.export时没指定dynamic_axes导出的ONNX会把H和W锁死成固定值。比如你训练时用128x128导出后只能处理128x128的图换个尺寸直接报错。正确的写法是这样# 这里踩过坑dynamic_axes必须显式声明input_names[input]output_names[output]dynamic_axes{input:{2:height,3:width},output:{2:height,3:width}}torch.onnx.export(model,dummy_input,model.onnx,input_namesinput_names,output_namesoutput_names,dynamic_axesdynamic_axes,opset_version11,# 别用太新的opset有些移动端框架不支持do_constant_foldingTrue)注意那个opset_version我习惯用11。opset 13以上有些算子比如Resize的坐标变换模式在NCNN里会炸后面移植时你会感谢这个选择。导出后一定要用onnxruntime验证一遍importonnxruntimeasort sessionort.InferenceSession(model.onnx)# 别这样写直接拿训练时的预处理数据# 要用实际推理时的输入格式比如BGR转RGB、归一化参数等input_datanp.random.randn(1,3,256,256).astype(np.float32)outputssession.run(None,{input:input_data})如果输出和PyTorch结果差超过1e-3八成是Batch Normalization或Instance Normalization的融合问题。超分模型里常见的nn.PixelShuffle在ONNX里会被拆成ReshapeTranspose这一步容易丢精度建议导出前用torch.jit.trace先跑一遍看计算图有没有异常节点。二、TensorRT加速FP16推理的精度陷阱ONNX导好了接下来用TensorRT做推理加速。这一步的收益肉眼可见——一个EDSR模型在RTX 3060上FP32推理要15ms转成FP16直接降到5ms。但代价是精度波动。我遇到过最离谱的情况一个4倍超分模型FP16推理后输出图像出现周期性条纹像摩尔纹一样。查了两天才发现是Conv2d的权重在FP16下溢出某些通道的激活值太大被截断了。解决方案是逐层精度校准。别一股脑全转FP16用TensorRT的Int8Calibrator思路做逐层分析# 这里踩过坑不要直接builder.create_network(1int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))# 超分模型通常有动态batch用EXPLICIT_BATCH模式networkbuilder.create_network(1int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))parsertrt.OnnxParser(network,logger)parser.parse_from_file(model.onnx)# 关键步骤逐层设置精度forlayerinnetwork:iflayer.typetrt.LayerType.CONVOLUTION:# 卷积层用FP16但输出通道数大的层比如最后一层用FP32iflayer.get_output(0).shape[1]64:layer.precisiontrt.float32else:layer.precisiontrt.float16更省事的做法是用trtexec工具跑一遍校准它会自动标记哪些层不适合低精度。命令如下trtexec--onnxmodel.onnx--fp16--saveEnginemodel_fp16.engine--calibcalibration_data校准数据最好用真实场景的图片别用ImageNet的随机裁剪。超分模型对纹理细节敏感校准集如果全是平滑区域FP16推理时高频细节会糊掉。三、NCNN移植从ncnn2onnx到param/bin文件移动端部署NCNN是绕不开的。但NCNN对ONNX的支持并不完美尤其是超分模型里常见的PixelShuffle和DepthToSpace算子。第一步用onnx2ncnn工具转换./onnx2ncnn model.onnx model.param model.bin如果报错说“unsupported operator”别慌。NCNN的算子支持列表里Resize、Conv、Relu这些基本都有但PixelShuffle可能被映射成ShuffleChannel或Reshape。我遇到过PixelShuffle被拆成四个Reshape加一个Permute效率极低。手动优化方案在PyTorch里把PixelShuffle替换成DepthToSpaceNCNN对后者支持更好。或者干脆在导出ONNX前把PixelShuffle用torch.nn.functional.pixel_shuffle显式写出来别用nn.PixelShuffle模块这样ONNX的算子映射更直接。转换成功后用ncnnoptimize做图优化./ncnnoptimize model.param model.bin model_opt.param model_opt.bin0那个数字0代表FP321代表FP16。移动端建议用FP16但要注意高通骁龙8 Gen 1以上才支持FP16加速老芯片会回退到FP32速度反而更慢。四、移动端推理内存与速度的博弈在手机上跑超分模型最头疼的不是计算量而是内存带宽。一个4K输入的超分模型中间特征图动辄几百MB手机那点内存根本扛不住。我的经验是分块推理。把输入图切成256x256的块每块单独推理最后拼回去。但直接拼会有接缝需要用重叠区域做融合# 别这样写直接硬拼# 要用overlap加权平均overlap16# 重叠像素数weightnp.ones((patch_size,patch_size))# 边缘部分渐变权重避免接缝weight[:overlap,:]*np.linspace(0,1,overlap)[:,None]weight[-overlap:,:]*np.linspace(1,0,overlap)[:,None]weight[:,:overlap]*np.linspace(0,1,overlap)[None,:]weight[:,-overlap:]*np.linspace(1,0,overlap)[None,:]这个技巧在NCNN上实测有效接缝几乎不可见。但代价是推理次数增加比如一张1080p的图要切40多块每块推理5ms总耗时200ms勉强能接受。另一个优化点是输入输出复用内存。NCNN的Extractor支持set_input和extract但每次调用都会申请新内存。用ncnn::Mat的clone_from方法复用缓冲区能省下30%的内存分配开销。五、个人经验别信文档信你的log最后说点实在的。部署这件事80%的时间在debug20%的时间在写代码。我踩过的坑列出来能写一本书ONNX导出时opset_version选11别选13以上NCNN的Resize算子会炸TensorRT的FP16推理一定要用校准集别偷懒NCNN的PixelShuffle支持不好提前在PyTorch里换成DepthToSpace移动端推理分块重叠融合是王道别想着一次跑完整图所有精度验证用PSNR和SSIM说话别只看肉眼如果你现在正在部署一个超分模型建议先从最简单的EDSR或SRCNN开始把整个流程跑通再换复杂模型。别一上来就搞ESRGAN那个生成器的跳跃连接在ONNX里会变成一团乱麻debug到你怀疑人生。还有记得在代码里加# TODO: 这里踩过坑的注释三个月后你回来看会感谢自己的。