最近在部署一个BERT-like的小模型(大概110M参数),用PyTorch动态图推理单条样本大概12ms,想着转成ONNX用TensorRT加速一下。结果导出很顺利,但用onnxruntime-gpu跑下来居然要18ms,反而慢了50%?我确认了输入输出都是动态shape,也开了graph optimization level,甚至试了固定seq_len到128,依然没改善。查了下网上说可能是算子不支持导致fallback到CPU,但看profiler又显示在CUDA上。有没有大佬遇到过类似情况?是TensorRT不擅长处理Transformer结构,还是我动态shape的设置有问题?或者量化INT8才是正解?求指点,谢谢!
楼主
2026-08-05
PyTorch转ONNX后推理速度反而变慢,是姿势不对还是框架问题?
请 登录 后发表回复
全部回复
共 103 条
2楼
27天前
小模型动态shape下ORT未必比PyTorch快,试试固定shape跑一次看是不是shape推断在拖后腿。
3楼
23天前
你这个问题挺典型的,110M的BERT类模型转ONNX确实容易踩坑。onnxruntime的CUDA EP对MultiHeadAttention的融合不如TensorRT彻底,而且dynamic shape下很多kernel会走通用实现,性能掉一截很正常。建议先试下用TRT的EP跑(onnxruntime里面直接挂TensorrtExecutionProvider),或者干脆用polygraphy看看哪些节点被优化掉了。我之前也遇到过,固定shape后反而更慢,因为图优化把reshape给拆了,你可以对比下onnx和torch的profiler看瓶颈到底在哪一层。
4楼
23天前
ONNX Runtime的CUDA EP默认对Transformer里的Attention和LayerNorm融合得比较保守,没TensorRT那么激进,所以单看ORT不一定比PyTorch快。你可以试试拿onnxruntime的TensorRT EP跑,或者直接上trtexec把ONNX转成engine,通常能压到8ms以内。另外动态shape下ORT会频繁重新分配显存,固定seq_len反而可能触发不同kernel,建议把profile打开看看具体哪个节点耗时长。