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])