pyTorch onnx小试《一》

环境安装（conda内）

conda create -n onnx python=3.7 -y
conda activate onnx
conda install -c conda-forge protobuf numpy
pip install onnx
conda install pytorch torchvision

pyTorch版本AlexNet模型转换到onnx

import torch
import torchvision

dummy_input = torch.randn(10, 3, 224, 224, device='cuda') #定义输入的数据类型(B,C,H,W)为(10,3,224,224)
model = torchvision.models.alexnet(pretrained=True).cuda() #读入torchvision中已经训练好的alxnet模型
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]
torch.onnx.export(model, dummy_input, "alexnet.onnx", keep_initializers_as_inputs=True, verbose=True, input_names=input_names, output_names=output_names)

在onnx中运行:

import onnx
# Load the ONNX model
model = onnx.load("alexnet.onnx")

# Check that the IR is well formed
onnx.checker.check_model(model)

# Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)

import caffe2.python.onnx.backend as backend
import numpy as np

rep = backend.prepare(model, device="CUDA:0") # or "CPU"
# For the Caffe2 backend:
#     rep.predict_net is the Caffe2 protobuf for the network
#     rep.workspace is the Caffe2 workspace for the network
#       (see the class caffe2.python.onnx.backend.Workspace)
outputs = rep.run(np.random.randn(10, 3, 224, 224).astype(np.float32))
# To run networks with more than one input, pass a tuple
# rather than a single numpy ndarray.
print(outputs[0])

错误处理

1. ModuleNotFoundError: No module named ‘past’
   处理方法：

pip install future

2. IndexError: _Map_base::at
   处理方法：
   见 https://github.com/onnx/onnx/issues/2417

参考：
https://blog.youkuaiyun.com/u010946556/article/details/82734216?depth_1-utm_source=distribute.pc_relevant.none-task&utm_source=distribute.pc_relevant.none-task
https://github.com/nilmtk/nilmtk/issues/548
https://github.com/onnx/onnx/issues/2417