ONNX

ONNX官网链接

首先将自己训练出来的pth模型转化为ONNX模型。
这是我自己的，改为自己的即可

import torch
import torchvision.models as models

import json
import os
import datetime

import torch
from torch.utils import data
import numpy as np

import transforms
from model import HighResolutionNet
from my_dataset_coco import CocoKeypoint
from train_utils import train_eval_utils as utils
model=HighResolutionNet(32,3)
checkpoint = torch.load('./save_weights/model-42.pth', map_location='cpu')
model.load_state_dict(checkpoint['model'])

# 创建虚拟输入数据
dummy_input = torch.randn(1, 3, 512, 512)  # 示例输入形状

# 导出模型到 ONNX 格式
output_file = "model1.onnx"
torch.onnx.export(model, dummy_input, output_file)

print("模型已成功转换为 ONNX 格式并保存为:", output_file)

之后将在网站中打开该onnx文件 即可看到网络模型