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(angle + margin))
        denominator_last = torch.sum(torch.exp(s * torch.cos(angle)), dim=1, keepdim=True) - torch.exp(s * torch.cos(angle))
        output = torch.log(numerator / (numerator + denominator_last))
        return output

class FaceNet(nn.Module):
    def __init__(self):
        super(FaceNet, self).__init__()
        self.sub_net = models.densenet169()
        in_feature = self.sub_net.classifier.in_features

        self.sub_net.classifier = nn.Sequential(
            nn.Linear(in_feature, 1000),
            nn.ReLU(inplace=True),
            nn.Linear(1000, 512))

        self.arc_softmax = Arcloss(512, 1855)

    def forward(self, x):
        feature = self.sub_net(x)
        return feature, self.arc_softmax(feature)

    def encode(self, x):
        return self.sub_net(x)

# 转onnx
def dsnet_to_onnx():
    # 初始化模型
    model = FaceNet()
    model.load_state_dict(torch.load(r'param_wd\dsnt169_1855_4.pth'))
    model.eval()

    # 创建一个虚拟输入张量   (在这里有些博主会把虚拟张量也进行数据预处理，实际上不需要。因为：
    # 在将模型转换为 ONNX 格式时，定义的虚拟输入张量 dummy_input = torch.randn(1, 3, 240, 240) 只是为了确保模型的输入形状是正确的。预处理是在实际推理时进行的，因此在转换过程中不需要对虚拟输入进行预处理。)
    dummy_input = torch.randn(1, 3, 240, 240)

    # 导出模型
    torch.onnx.export(model,
                      dummy_input,
                      "facenet.onnx",            # 保存的模型名称
                      opset_version=10,
                      do_constant_folding=True,  # 是否执行常量折叠优化
                      input_names=["input"],  # 输入名
                      # output_names=['output'] # 如果你的网络的输出只有一个（比如单纯的分类网络，输出就是类别，输出名就用这个就行）
                      output_names=['feature', 'output'],  # 输出名 （因为我的网络的输出是两个值，特征与类别，所以这里的输出名定义的两个）
                      dynamic_axes={"input": {0: "batch_size"},  # 批处理变量
                                    "output": {0: "batch_size"}})

# 查看网络结构（可选）
def see_net():
    modelData = "facenet.onnx"
    netron.start(modelData)


def resize_and_pad_image(image_path_or_imdata, target_size=(240, 240), background_color=(0, 0, 0), is_path=True):
    # 打开图片
    if is_path:
        img = Image.open(image_path_or_imdata)
    else:
        img = image_path_or_imdata
    if img.mode != "RGB":
        img = img.convert("RGB")
    # 获取图片的原始尺寸
    width, height = img.size

    # 计算缩放比例
    ratio = min(target_size[0] / width, target_size[1] / height)
    new_width = int(width * ratio)
    new_height = int(height * ratio)

    # 缩放图片
    img = img.resize((new_width, new_height), Image.LANCZOS)

    # 创建新的背景图
    background = Image.new('RGB', target_size, background_color)

    # 计算图片放置的位置
    x = (target_size[0] - new_width) // 2
    y = (target_size[1] - new_height) // 2

    # 将缩放后的图片粘贴到背景图上
    background.paste(img, (x, y))

    return background


def onnx_infer(img):

    # 加载 ONNX 模型 (优先使用gpu)
    providers = [
        ('CUDAExecutionProvider', {
            'device_id': 0  # 使用 GPU 0
        }),
        'CPUExecutionProvider'  # 后备提供者
    ]
    ort_session = onnxruntime.InferenceSession("facenet.onnx", providers=providers)
    print("加载onnx模型成功")
    # 查看会话中使用的执行提供者
    providers = ort_session.get_providers()
    print("使用的执行提供者:", providers)
    tf = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
    img = resize_and_pad_image(img, is_path=True)
    person_picture = tf(img).numpy()                         # 将 Tensor 转换为 NumPy 数组
    person_picture = np.expand_dims(person_picture, axis=0)  # 添加批次维度

    # 运行推理
    outputs = ort_session.run(None, {'input':  person_picture})

    # 输出特征和分类结果
    feature, output = outputs

    # print("Feature:", feature)
    # print("Output:", output)
    return feature, output


# 计算欧氏距离或余弦相似度
def compute_similarity(feature1, feature2):
    return 1 - cosine(feature1.flatten(), feature2.flatten())  # 计算余弦相似度


if __name__ == '__main__':
    while True:
        feature1, cls1 = onnx_infer(img='0c.jpg')
        feature2, cls2 = onnx_infer(img='17cjh.jpg')
        # 计算相似度
        similarity = compute_similarity(feature1, feature2)
        print(f"Similarity: {similarity:.4f}")