pytorch模型转onnx(转载)

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

1.torch模型转onnx

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

1. import torchvision

2.  

3. dummy_input = torch.randn(10, 3, 224, 224, device='cuda')

4. model = torchvision.models.alexnet(pretrained=True).cuda()

5.  

6. # Providing input and output names sets the display names for values

7. # within the model's graph. Setting these does not change the semantics

8. # of the graph; it is only for readability.

9. #

10. # The inputs to the network consist of the flat list of inputs (i.e.

11. # the values you would pass to the forward() method) followed by the

12. # flat list of parameters. You can partially specify names, i.e. provide

13. # a list here shorter than the number of inputs to the model, and we will

14. # only set that subset of names, starting from the beginning.

15. input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]

16. output_names = [ "output1" ]

17.  

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

19.  

20. # not assign the input name and output name

21. #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进行推理

1. # ...continuing from above

2. import onnxruntime as ort

3.  

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

5.  

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

7.  

8. print(outputs[0])