pytorch 分类模型转onnx并推理

1. 先安装包： pip install onnx   onnxruntime  netron  onnxruntime-gpu  -i https://pypi.tuna.tsinghua.edu.cn/simple

2.转换推理代码

import torch
import torch.onnx
from torchvision import models
import torch.nn as nn
import torch.nn.functional as F
import netron
import os, sys
sys.path.append(os.getcwd())
import onnxruntime
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
from scipy.spatial.distance import cosine

class Arcloss(nn.Module):
    def __init__(self, feature_dim=512, cls_dim=20):
        super(Arcloss, self).__init__()
        self.W = nn.Parameter(torch.randn(feature_dim, cls_dim), requires_grad=True)

    def forward(self, features, margin=1, s=10):
        x = F.normalize(features, dim=1)
        w = F.normalize(self.W, dim=0)
        cosa = torch.matmul(x, w) / s
        angle = torch.acos(torch.clamp(cosa, -1.0 + 1e-7, 1.0 - 1e-7))
        numerator = torch.exp(s * torch.cos(ang