pytorch模型转onnx

pytorch模型转其他的格式一般需要先转成通用的格式，onnx就是一种通用的格式，下面介绍转onnx的方法

1.torch模型转onnx

  直接通过torch官网：https://pytorch.org/docs/stable/onnx.html，提供的方法转onnx

import torchvision

dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
model = torchvision.models.alexnet(pretrained=True).cuda()

# Providing input and output names sets the display names for values
# within the model's graph. Setting these does not change the semantics
# of the graph; it is only for readability.
#
# The inputs to the network consist of the flat list of inputs (i.e.
# the values you would pass to the forward() method) followed by the
# flat list of parameters. You can partially specify names, i.e. provide
# a list here shorter than the number of inputs to the model, and we will
# only set that subset of names, starting from the beginning.
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]

torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)

# not assign the input name and output name
#torch.onnx.export(model, dummy_input, "alexnet.onnx")

转换的过程中容易出问题，一般是torch和torchvision的版本导致，下面的版本验证没有问题

  pip install torch==1.2.0 -i  https://pypi.tuna.tsinghua.edu.cn/simple

  pip install torchvision=0.4.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

2. onnx 推理

  pip install onnx==1.7.0 -i  https://pypi.tuna.tsinghua.edu.cn/simple

  pip install onnxruntime-gpu=1.4.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

  安装onnx，然后使用onnxruntime进行推理

# ...continuing from above
import onnxruntime as ort

ort_session = ort.InferenceSession('alexnet.onnx')

outputs = ort_session.run(None, {'actual_input_1': np.random.randn(10, 3, 224, 224).astype(np.float32)})

print(outputs[0])