使用 ML.NET 启动模型
下载分类文件标签文件:
https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt
通过 Pytorch ,将 pth 格式模型转换为 onnx 格式模型。
from torchvision import models, datasets, transforms as T
import torch
from PIL import Image
import numpy as np
resnet50 = models.resnet50(pretrained=True)
# Read the categories
with open("imagenet_classes.txt", "r") as f:
categories = [s.strip() for s in f.readlines()]
# Export the model to ONNX
image_height = 224
image_width = 224
x = torch.randn(1, 3, image_height, image_width, requires_grad=True)
torch_out = resnet50(x)
torch.onnx.export(resnet50, # model being run
x, # model input (or a tuple for multiple inputs)
"resnet50.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=12, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['output']) # the model's output names
启动程序后得到 resnet50.onnx 文件。