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