Read the categories

使用 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 文件。