Skip to main content

如何将Torch模型转换到onnx

备注

在我们的所有已知实践中，通过静态onnx组合, 动态onnx，预生成tensorrt模型等途径，torchpipe能完整替代 torch2trt.

torch转onnx​

框架优先支持动态 batch 或者 batchsize==1 的静态 batch。实际中，有些模型无法转为动态尺度，或者比较容易出错， 我们也支持同时加载多个不同静态batchsize的模型，去模拟动态尺度。以下说明主要针对导出动态 batchsize 模型。

动态batch的导出
  • 以下操作导致动态batch不可用： x.view(int(x.size(0)), -1). 需要检查模型文件中是否存在将batch维度写死的情况，比如：x.view(int(x.size(0)), -1, 1, 1)，x.reshape(int(x.size(0)), -1, 1, 1)等，这可能会导致转换onnx后动态batch出现问题。注意，在Transformer-like的网络中，batch维度不一定在第0维度。
  • batch维度指定为动态大小时，低版本tensorrt对此处理能力弱一些，冗余算子多一些。比如对于 x.view(x.size(0), -1)，会在onnx中引入Gather等算子来计算x的第一个维度。可修改为 x = x.view(-1, int(x.size(1)*x.size(2)*x.size(3))) 或者 x = torch.flatten(x, 1)。此项非必需。
  • 对于部分模型（tensorrt8.5.1, lstm 和 transformer），batch维度和非batch维度同时动态时，可能消耗更多资源 ：
  • 对于layerNorm层以及动态batch的Transformer-like的网络，推荐使用opset>=17, tensorrt>=8.6.1
# batch和非batch同时动态，需要9ms（推理输入大小为optShapes=input:1x1000x80,mask:1x1x1000）:
/opt/tensorrt/bin/trtexec --onnx=test_fp32.onnx --shapes=input:1x1000x80,mask:1x1x1000 --workspace=64000 \
    --minShapes=input:1x20x80,mask:1x1x20 \
    --optShapes=input:1x1000x80,mask:1x1x1000 \
    --maxShapes=input:4x2000x80,mask:4x1x2000 
    

# 固定batchsize==1，只需要4.6ms：
/opt/tensorrt/bin/trtexec --onnx=test_fp32.onnx --shapes=input:1x1000x80,mask:1x1x1000 --workspace=64000 \
    --minShapes=input:1x20x80,mask:1x1x20 \
    --optShapes=input:1x1000x80,mask:1x1x1000 \
    --maxShapes=input:1x2000x80,mask:1x1x2000

此时推荐只将其中一个维度离散化。

最佳实践

修改完网络后，可以利用下面代码，将pytorch模型转换为onnx模型。

x = torch.randn(1,3,224,224).cuda()
out_size = 1
in_out={"input":{0:"batch_size"}}
for i in range(out_size):
    in_out[f"output_{i}"] = {0:"batch_size"}

torch_model.cuda().eval()
torch.onnx.export(torch_model,
                    x,
                    onnx_save_path,
                    opset_version=17,
                    do_constant_folding=True,   
                    input_names=["input"],      # 输入名
                    output_names=[f"output_{i}" for i in range(out_size)],  # 输出名
                    dynamic_axes=in_out)

import onnx
from onnxsim import onnx_simplifier

onnx_model = onnx.load(onnx_save_path)
onnx_model = onnx.shape_inference.infer_shapes(onnx_model)
model_simp, check = onnx_simplifier.simplify(onnx_model, check_n = 0)
onnx.save(model_simp, onnx_save_path)
为了方便，针对这步我们提供了torchpipe.utils.models.onnx_export小工具(0.3.2b1开始生效)
  • 该工具可以实现将PyTorch模型转换为ONNX模型并保存到本地。
  • 支持动态batch，并且自带onnx-simplify优化。
def onnx_export(model: Union[torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction], onnx_path, input = None, opset = 17):
参数
  • model - PyTorch模型。
  • onnx_path - ONNX模型保存路径。
  • input - 模型输入，如果不设置默认为torch.randn(1,3,224,224)。
  • opset ONNX的opset。
示例代码
import os, tempfile
from torchvision import models
import torch
import torchpipe

## export onnx
m = models.resnet50(weights=None).eval()
onnx_path = os.path.join(tempfile.gettempdir(), f"resnet50.onnx")
torchpipe.utils.models.onnx_export(m, onnx_path, torch.randn(1, 3, 224, 224), opset=17)

转换失败说明​

torch转onnx经常遇到转换失败的情况。可采取的方法有：

  • 动态维度保持动态， 比如对于yolox：
x = x.view(int(x.size(0)), -1, 1, 1)
# 改为
x = x.flatten(1).unsqueeze(2).unsqueeze(2)

x = x.view(int(x.size(0)), -1)
# 改为
x = x.view(-1, int(x.size(1)*x.size(2)*x.size(3)))
  • bool值改为float:
tgt_padding_mask = (tgt_in == self.eos_id)
# 改为
tgt_padding_mask = (tgt_in == self.eos_id).float()
  • 版本原因, 尝试不同版本：
    • 尽可能使用最新版，比如采用 onnx opset >= 14 和 tensorrt >= 8.2
    • 对于tensorrt7，推荐onnx版本1.9.0, onnx opset = 11
  • 尝试使用trtexec进行模型转换：
    
/opt/tensorrt/bin/trtexec --onnx=folded.onnx 
--workspace=64000 \
--minShapes=input:1x20x80,mask:1x1x20 \
--optShapes=input:1x1000x80,mask:1x1x1000 \
--maxShapes=input:1x2000x80,mask:1x1x2000 \
--saveEngine=test_fp32.trt

onnx 相关工具​

onnx-simplify​

简化模型结构的工具

pip install onnx onnxsim
onnxsim input.onnx output.onnx

netron​

用于可视化onnx模型的工具。

运行 pip install netron 和 netron [FILE] 或者 netron.start('[FILE]').

ONNX GraphSurgeon​

ONNX GraphSurgeon 是tensorrt官方发布的一款用于修改onnx结构的工具。

Polygraphy​

nvidia官方用于测试tensorrt或者onnx的工具。提供模型转换功能，对于fp16精度损失可进行调试，指定层不使用fp16.

参考连接​