ONNX 静态图转换为动态图

模型转换

ONNX 模型从动态图变为静态图是非常容易的，使用 onnxsim --overwrite-input-shape input_name:dim0,dim1,dim2 model.onnx 即可，onnxsim 内部会将动态维度转为静态维度，同时执行图优化。

但是如果只有静态图，想要变回动态图，则比较困难。一般来说，Batch 维度是最容易改变的，修改 Batch 维度可以一次推理 N 个输入，提高吞吐量。

使用 onnx2torch 可以解决这个问题。首先将 ONNX 模型转为 pytorch 模型，然后重新使用 torch.onnx.export 导出一个有动态维度的模型即可。代码如下：

import onnx
import torch
from onnx2torch import convert

# onnx -> torch
onnx_model_path = "/some/path/mobile_net_v2.onnx"
torch_model_1 = convert(onnx_model_path)

# torch -> onnx
export_path = "/some/path/mobile_net_v2_torch.onnx"
torch.onnx.export(
  torch_model_1,
  (torch.randn(1, 3, 224, 224),),
  export_path,
  opset_version=11,
  export_params=True,
  input_names=["input"],
  output_names=["output"],
  dynamic_axes={"input": {0: "B"}}
)
PY

注意：由于静态图之前可能被执行过图优化等操作，一部分节点会被合并或者删除，因此 onnx2torch 转换之后的模型不一定能够用于模型训练，只能用于推理。