pytorch训练结果 转 onnx

pytorch训练结果转可部署onnx

import torch
import torchvision

input_data = torch.randn(1, 3, 224, 224)
model = torchvision.models.resnet18(pretrained=True)
input_names = [ "input" ]
output_names = [ "output"]

torch.onnx.export(model,
input_data ,
"resnet.onnx",
verbose=True,
input_names=input_names,
output_names=output_names)

# check onnx model
onnx_model = onnx.load("resnet.onnx") # load onnx model
onnx.checker.check_model(onnx_model)

转自：https://blog.csdn.net/weixin_41798111/article/details/122574833?spm=1001.2014.3001.5501

标签：load,训练,onnx,pytorch,names,output,input,model
来源： https://www.cnblogs.com/student-gao/p/15871573.html