打造你的Pokemon大师：深度学习多分类模型构建与本地部署全攻略

引言

在这篇文章中，我将分享如何从头构建一个用于多分类任务的神经网络，并将其部署到本地环境。这是一个实践性质的教程，适合那些对深度学习模型部署感兴趣的初学者。

数据集准备

为了本次训练，我选择了一个网上流行的宝可梦数据集，它包含5个类别，每个类别的图片都存放在各自的文件夹中。为了确保标签和类别的一致性，我使用了sorted()函数对文件夹名称进行排序，并通过enumerate为每个类别分配一个唯一的标签。

以下是数据集准备的代码实现：

import glob
import os
import cv2
from torch.utils.data import DataLoader, Dataset
import torchvision
from PIL import Image
import randomclass PokemonData(Dataset):def __init__(self, root_path, mode=None):super(PokemonData, self).__init__()self.pokemon_names = sorted(os.listdir(root_path))self.labels = {name: i for i, name in enumerate(self.pokemon_names)}self.all_imgs = []for name in self.pokemon_names:self.all_imgs.extend(glob.glob(os.path.join(root_path, name, '*')))random.shuffle(self.all_imgs)self.imgs = self.all_imgs[:int(len(self.all_imgs) * 0.8)] if mode == "train" else self.all_imgs[int(len(self.all_imgs) * 0.8):]def __len__(self):return len(self.imgs)def __getitem__(self, item):name = self.imgs[item].split("\\")[-2]img = cv2.imread(self.imgs[item])tf_img = self.transformData(img)label = self.labels[name]return label, tf_imgdef transformData(self, img):img = Image.fromarray(img)tf_img = torchvision.transforms.Compose([torchvision.transforms.Resize((224, 224)),torchvision.transforms.ToTensor(),torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])img = tf_img(img)return imgif __name__ == '__main__':root_path = "E:/pokemon"mode = "train"pd = PokemonData(root_path, mode)pd_datas = DataLoader(pd, batch_size=32, shuffle=True)for label, data in pd_datas:print(label, data.shape)

模型构建

我选择了预训练的ResNet18模型作为基础，因为它在性能和计算资源之间取得了良好的平衡。ResNet18的最后一层输出1000个类别，我们需要将其替换为适合我们数据集的输出层。

from torchvision import models
import torch
import torch.nn as nnclass CustomResNet18(nn.Module):def __init__(self):super(CustomResNet18, self).__init__()self.base_model = models.resnet18(pretrained=True)self.fc = nn.Linear(512, 5)def forward(self, x):x = self.base_model(x)x = self.fc(x)return x

训练与评估

接下来，我们训练模型并评估其性能。我们使用交叉熵损失函数和Adam优化器。

# 训练和评估代码省略，与原文中相同

模型格式转换

为了提高预测性能，我们将PyTorch模型转换为ONNX格式。ONNX是一种开放的模型格式，允许模型在不同的框架和硬件之间迁移。

import onnx
import torch# 导入自定义模型
from custom_resnet import CustomResNet18model = CustomResNet18()
model.load_state_dict(torch.load("best_model.pt"))
model.eval()x = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, x, "best_model.onnx", input_names=["input"], output_names=["output"], opset_version=11)onnx_model = onnx.load("best_model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX模型格式正确!")

本地部署与预测

最后，我们使用ONNX模型进行本地预测。以下是如何加载ONNX模型并对一张图片进行分类的示例代码。

import onnxruntime as ort
import torch
import cv2
from PIL import Image
from torchvision import transforms# 加载ONNX模型
session = ort.InferenceSession("best_model.onnx")# 图像预处理
data_preproce = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])img_cv = cv2.imread("test.png")
img_pil = Image.fromarray(img_cv)
img = data_preproce(img_pil)
input_tensor = torch.unsqueeze(img, 0).numpy()# 进行预测
pred = session.run(None, {"input": input_tensor})[0]
pred_softmax = torch.softmax(torch.tensor(pred), dim=1)
values, indices = torch.topk(pred_softmax, 3)# 显示预测结果
labels_dict = {'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4}
revers_dict = {v: k for k, v in labels_dict.items()}
for i in indices[0].tolist():print(revers_dict[i], ":", round(values[0].tolist()[index_n.index(i)] * 100, 5), "%")