将模型转为onnx
选用pytorch框架,训练resnet18二分类,将二分类模型转为onnx模型
from typing_extensions import dataclass_transform
import torch
import torch.nn as nn
from torchvision import models, transforms
import cv2 as cv
import numpy as np
from PIL import Image
class_name = ['ant', 'bee']
# device = torch.device('cuda' if torch.cuda.is_available else 'cpu')
img = Image.open(r'D:\Projects\PythonProjects\ResnetClassifier\data\hymenoptera_data\val\ants\800px-Meat_eater_ant_qeen_excavating_hole.jpg')
img = img.convert('RGB')
img = img.resize((224, 224))
img = np.array(img, np.float32)
preprocess_img = img/255.0
img_data = np.expand_dims(np.transpose(preprocess_img, (2, 0, 1)), 0)
device = torch.device('cpu')
model = models.resnet18()
model_l = model.fc.in_features
model.fc = nn.Linear(model_l, 2)
model_x = model.to(device)
model_x.eval()
model_x.load_state_dict(torch.load("D:\\Projects\\PythonProjects\\Datasets\\resnet18_bee_and_ant.pth", map_location=torch.device('cpu')))
tensor_img = torch.from_numpy(img_data)
out = model_x(tensor_img)
_, pred = torch.max(out<