最近在把一个训练好的分割模型部署到服务器上，看教程都说要用ONNX转一下再上TensorRT。但我自己试了试，直接用PyTorch的torch.jit.script或者直接加载权重跑推理，速度也没差太多啊？而且转ONNX的时候还踩了不少坑，像动态尺寸、算子不支持这些，改了半天才跑通。想问问大家，在实际生产环境里，用PyTorch原生做推理到底卡在哪？是显存占用、多线程并发，还是说TensorRT的优化幅度真的很大，只是我模型太小没感觉出来？有没有大佬说说自己踩过的坑，让我少走点弯路。
楼主
4天前
PyTorch的autograd明明很方便，为什么还要转成ONNX部署？
请 登录 后发表回复
全部回复
共 1 条
2楼
6小时前
你这个感受其实挺正常的，模型小的时候确实看不出太大差距，尤其是batch size固定、输入尺寸不变的情况下，torch.jit.script跑得也不慢。但真到生产环境，问题往往不是单次推理速度，而是并发和显存。PyTorch原生推理每个请求都可能带一堆Python解释器开销和GIL限制，多线程一上来就炸，而TensorRT在kernel融合和fp16/int8量化上的优化，大模型上能差出两三倍甚至更多。我之前部署过一个BERT类模型，PyTorch直接跑延迟大概40ms，转ONNX再上TRT直接降到12ms左右，而且显存占用少了快一半。ONNX那些坑确实烦，动态shape、自定义算子、版本兼容，基本每个模型都要折腾一遍，但一旦跑通，后面推理服务稳定性和吞吐量提升很明显。另外PyTorch Serving也不是不能用，只是生态和优化深度跟TRT比还是差一截，尤其你要做多模型流水线或者边缘部署的时候。所以不是ONNX有多好，而是生产环境要的是可预测的延迟和资源占用，这点原生PyTorch确实吃亏。