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

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

引用

CSDN

1.

https://m.blog.csdn.net/m0_73388849/article/details/143532734

本文将详细介绍如何从头构建一个用于多分类任务的神经网络，并将其部署到本地环境。通过这篇文章，你将学习到数据集准备、模型构建、训练评估、模型格式转换以及本地部署的完整流程。

数据集准备

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

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

import glob
import os
import cv2
from torch.utils.data import DataLoader, Dataset
import torchvision
from PIL import Image
import random

class 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_img

    def 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 img

if __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 nn

class 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 CustomResNet18

model = 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), "%")

通过这篇文章，不仅学习了如何构建和训练一个多分类神经网络，还了解了如何将其部署到本地环境并进行预测。希望这篇文章对你有所帮助！