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)