文章目录


好久没有自己写一个多分类的网络了，所以回顾一下，仅用于记录

一、数据集准备

​ 本次训练用的数据集是在网上找来的宝可梦数据集，具体的来源忘记了，一共是5个分类，每个文件夹都放着不同样式的精灵，如下图。
在这里插入图片描述

​ 这里有5个分类，所有每个分类要对应一个标签，标签从0到4。为了使得每次执行代码的时候，每个标签和每个分类都对应得上，我采用的是sorted()给每个文件夹的名称排好序，然后再通过enumerate给每个文件夹都赋上一个值，这个值就是它们的标签，其中文件夹的名字我将其设定为标签的名字。

​ 下面是代码的实现，每个文件夹所对应的标签为：

​ {‘bulbasaur’: 0, ‘charmander’: 1, ‘mewtwo’: 2, ‘pikachu’: 3, ‘squirtle’: 4}


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


class Pokmans_Data(Dataset):
    def __init__(self, root_path, mode=None):
        super(Pokmans_Data, self)
        self.pokman_names = sorted(os.listdir(root_path))
        self.labels = {}
        for i, name in enumerate(self.pokman_names):
            self.labels[name] = i
        self.all_imgs = []
        for name in self.pokman_names:
            self.all_imgs.extend(glob.glob(os.path.join(root_path, name, '*')))
        random.shuffle(self.all_imgs)
        if mode == "train":
            self.imgs = self.all_imgs[:int(len(self.all_imgs) * 0.8)]
        if mode == "test":
            self.imgs = self.all_imgs[int(len(self.all_imgs) * 0.8):]
        print(self.labels)

    def __len__(self):
        return len(self.imgs)

    def __getitem__(self, item):
        # "E:\pokeman\bulbasaur\00000000.png"
        name = self.imgs[item].split("\\")[-2] # bulbasaur
        img = cv2.imread(self.imgs[item])
        tf_img = self.tansfromData(img)
        label = self.labels[name] # 0
        return label, tf_img

    def tansfromData(self, img):
        img = Image.fromarray(img)
        tf_img = torchvision.transforms.Compose([
            # Resize一定一定要写元祖形式，不然只有一个位置变成224，很TM坑
            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__':
    # 每个精灵对应的标签
    # {'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4}
    root_path = r"E:\pokeman"
    mode = "train"
    pkd = Pokmans_Data(root_path, mode)
    pkd_datas = DataLoader(pkd, batch_size=32, shuffle=True)
    for label, data in pkd_datas:
        print(label, "\n", data.shape)

    # tensor([1, 1, 4, 0, 2, 1, 4, 4, 3, 0, 0, 0, 3, 3, 4, 1, 0, 1, 0, 3, 2, 3, 0, 3,3, 4, 3, 2, 1, 1, 1, 4])
    # torch.Size([32, 3, 224, 224])

    # for label, tf_img in pkd:
    #     print(label, tf_img.shape)
    # 0 torch.Size([3, 224, 224])
    # 3 torch.Size([3, 224, 224])
二、模型的搭建

​ 本次的模型我采用的resnet18，不用32的原因是我的笔记本电脑太垃圾了，1050的显卡，到时候跑起来太慢了，而且18的效果也很好了。

​ resnet18最后一层的输出是1000个类别，所以最后一层我们要自己写。在写之前我们需要弄明白这resnet18在接收到224x224的图片后它的输出是什么，我们才知道最后一层的输入该怎么写

from torchvision import models
import torch
import torch.nn as nn

res18=models.resnet18(pretrained=True)
a=torch.randn(1,3,224,224)
print(res18(a).shape) # torch.Size([1, 1000]) 一张图有1000个类别

new18=nn.Sequential(*list(res18.children())[:-1],
                    nn.Flatten())

# CHW拉伸层一维后，输出512，所有最后一层线性层第一个维度写512
print(new18(a).shape) # torch.Size([1, 512])

​ 根据上面的代码我们已经知道输出的类别数量以及去除最后一层后，它的输出是torch.Size([1, 512])，所有我们就可以搭建网络并训练，保存pt模型

import torch
from torchvision import models
import torch.nn as nn
import torch.optim as optim
from pokemans import Pokmans_Data
from torch.utils.data import DataLoader
import warnings
import matplotlib.pyplot as plt

warnings.filterwarnings("ignore", category=UserWarning)


class MyResnet18(nn.Module):
    def __init__(self):
        super(MyResnet18, self).__init__()
        self.resnet18 = models.resnet18(pretrained=True)
        self.net = nn.Sequential(*list(self.resnet18.children())[:-1],
                                 nn.Flatten(),
                                 nn.Linear(512, 5))

    def forward(self, x):
        return self.net(x)


def evalut(model, dataLoader):
    model.eval()  # 转出评估模式
    acc = 0
    for labels, datas in dataLoader:
        labels, datas = labels.to("cuda"), datas.to("cuda")
        with torch.no_grad():
            # 模型的输出是torch.Size([32, 5])，32张图片，每张图片有5个分类，所以输出是5
            pred = model(datas).argmax(dim=1)
        acc += torch.eq(pred, labels).sum().float().item()
    return acc / len(dataLoader.dataset)


if __name__ == '__main__':

    epochs = 10
    model = MyResnet18()
    # a=torch.randn((32,3,224,224))
    # print(model(a).shape)#torch.Size([32, 5])
    criterion = nn.CrossEntropyLoss()
    model = model.to("cuda")
    criterion = criterion.to("cuda")
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    trainDatas = Pokmans_Data(r"E:\pokeman", mode="train")
    testDatas = Pokmans_Data(r"E:\pokeman", mode="test")
    trainDataLoader = DataLoader(trainDatas, batch_size=32, shuffle=True)
    testDataLoader = DataLoader(testDatas, batch_size=32, shuffle=True)
    best_acc, best_epoch = 0, 0
    epoch_list = []
    loss_list = []

    for epoch in range(epochs):
        for i, (labels, datas) in enumerate(trainDataLoader):
            # 将数据集放到cuda中训练
            labels, datas = labels.to("cuda"), datas.to("cuda")

            model.train()
            loss = criterion(model(datas), labels) # 损失函数
            optimizer.zero_grad() # 清空梯度
            loss.backward() # 反向传播
            optimizer.step() # 优化器更新参数
        print("epoch:", epoch, "loss:", loss.item())

        # 根据测试集的输出，保存最好的pt模型
        if epoch % 1 == 0:
            val_acc = evalut(model, testDataLoader)  # 测试集的acc
            if val_acc > best_acc:
                best_acc, best_epoch = val_acc, epoch
                # 会保存模型的结构（架构）以及模型的参数（权重和偏差）
                torch.save(model, "best_model1.pt")
                # 只保存模型的参数，不包括模型的结构
                torch.save(model.state_dict(), "best_model2.pt")

        # 绘制每一轮loss的变换的折线图。
        epoch_list.append(epoch)
        loss_list.append(loss.item())
        best_index = epoch_list.index(best_epoch)
        best_loss = loss_list[best_index]
        plt.plot(epoch_list, loss_list, color="red", linewidth=2)
        plt.title("loss change")
        plt.xlabel("epoch")
        plt.ylabel("loss")
        plt.savefig("loss_change.png")

    print("best_acc:", best_acc, "best_epoch:", best_epoch, "loss:", best_loss, "val_acc:", val_acc)
    plt.show()
    plt.close()

​ 损失函数的变化曲线如下图

在这里插入图片描述

​ 文件夹目前的情况如下，其中有几个py文件是目前还没有用到的

在这里插入图片描述

如果有多张显卡，使用多张显卡来训练的话采用：net=nn.DataParallel(model,device_ids=[0,1])，那么main函数以下的代码如下：

多GPU的训练net=nn.DataParallel(model,device_ids=[0,1]),需要调整batchsize的大小和学习率，batchsize会影响训练的时间，lr会影响精度

如果不调整的话有可能test acc 会变低，并且运行的时间只缩短了一点点，而有时候，样本的多样性决定了batchsize的大小。在这里我们先忽略掉多样性的问题，就多GPU的问题讨论。这里就是因为通讯的开销大于计算的开销，单GPU的时候是batchsize=32，而2个GPU的时候可以batchsize=32*2，lr可以大一点，所有需要自己去调整实验实验，下方的代码只改了个API，其余都没有改，有需要的可以自己修改下

#上面的代码都不需要修改
if __name__ == '__main__':
	epochs = 10
    model = MyResnet18().to("cuda")
    # 有多少张显卡就写多少个
    net=nn.DataParallel(model,device_ids=[0,1])
    # a=torch.randn((32,3,224,224))
    # print(net(a).shape)#torch.Size([32, 5])
    
    criterion = nn.CrossEntropyLoss()
    criterion = criterion.to("cuda")
    optimizer = optim.Adam(net.parameters(), lr=0.001)
    trainDatas = Pokmans_Data(r"E:\pokeman", mode="train")
    testDatas = Pokmans_Data(r"E:\pokeman", mode="test")
    trainDataLoader = DataLoader(trainDatas, batch_size=32, shuffle=True)
    testDataLoader = DataLoader(testDatas, batch_size=32, shuffle=True)
    best_acc, best_epoch = 0, 0
    epoch_list = []
    loss_list = []

    for epoch in range(epochs):
        for i, (labels, datas) in enumerate(trainDataLoader):
            # 将数据集放到cuda中训练
            labels, datas = labels.to("cuda"), datas.to("cuda")

            net.train()
            loss = criterion(net(datas), labels)  # 损失函数
            optimizer.zero_grad()  # 清空梯度
            loss.backward()  # 反向传播
            optimizer.step()  # 优化器更新参数
        print("epoch:", epoch, "loss:", loss.item())

        # 根据测试集的输出，保存最好的pt模型
        if epoch % 1 == 0:
            val_acc = evalut(net, testDataLoader)  # 测试集的acc
            if val_acc > best_acc:
                best_acc, best_epoch = val_acc, epoch
                # torch.save(model, "best_model1.pt")
                # torch.save(model.state_dict(), "best_model2.pt")

        # 绘制每一轮loss的变换的折线图。
        epoch_list.append(epoch)
        loss_list.append(loss.item())
        best_index = epoch_list.index(best_epoch)
        best_loss = loss_list[best_index]
        plt.plot(epoch_list, loss_list, color="red", linewidth=2)
        plt.title("loss change")
        plt.xlabel("epoch")
        plt.ylabel("loss")
        # plt.savefig("loss_change.png")

    print("best_acc:", best_acc, "best_epoch:", best_epoch, "loss:", best_loss, "val_acc:", val_acc)
    plt.show()
    plt.close()
三、模型格式转换

​ 我们将pt转出onnx类型，为什么呢？因为pt在预测的时候它的fps很低，onnx将某些层进行了合并剪枝操作。并且onnx一种通用的模型交换格式，可以用于将模型从一个深度学习框架转换到另一个深度学习框架或硬件平台。

格式转换我是重新创建了一个新的py文件，并且这个格式转换会根据自己保存模型的方式有关，所有特别容易报错，下面我将两种保存模型的方式，并转成onnx的方式一一列举

两种方式都是需要去我们的模型搭建里，将我们的模型导入进来，注意这个模型一定要放在main主函数的上面，就像我上面自己写的搭建模型那里一样。。不然导入的时候，它会直接把整个py文件给执行了，然而我们只想导入一个模型的架构而已，并不要求它要执行一整个文件

方式1.torch.save(model, “best_model1.pt”)

这种方式一定要导入自己搭建的模型，不然会出现类似的 报错信息：AttributeError: Can’t get attribute ‘MyResnet18’ on <module ‘main’ from ‘E:/my_net/transfrom_pt_onnx.py’>

import onnx
import torch

# 导入自己定义的模型(第二部分模型的搭建那里)，
from my_resnet import MyResnet18

devices=torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load("best_model1.pt")
model.eval()
x= torch.randn(1, 3, 224, 224).to(devices)
# 为什么export需要一个输入呢？
# 因为转出onnx本质就是一种语言的翻译，从一种格式转出另一种。torch提供了一种追踪机制，给一个输入，让模型在执行一遍
# 将这个输入对应的计算图记录下来并保存成onnx
with torch.no_grad():
    ort_model = torch.onnx.export(
        model,  # 要转换的模型
        x,  # 任意一组输入(记得是float类型)
        "best_model1.onnx",  # 导出的文件
        input_names=["input"],  # 输入tensor的名称
        output_names=["output"],  # 输出tensor的名称
        opset_version=11)  # onnx算子集版本号

# 读取onnx模型
onnx_model=onnx.load("best_model1.onnx")
# # 检查模型格式是否正确
onnx.checker.check_model(onnx_model)
print("onnx模型格式正确!")
方式2.torch.save(model.state_dict(), “best_model2.pt”)

这种方式，在导入参数时候，一定一定一定一定一定不要赋值，不要赋值！！

import onnx
import torch
from my_resnet import MyResnet18

# 有GPU就用GPU，没有就用CPU
devices=torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyResnet18()
model.load_state_dict(torch.load("best_model2.pt")) # 这里一定不要赋值给model，否则会报错
model=model.to(devices)

model.eval()
x= torch.randn(1, 3, 224, 224).to(devices)
with torch.no_grad():
    ort_model = torch.onnx.export(
        model,  # 要转换的模型
        x,  # 任意一组输入(记得是float类型)
        "best_model2.onnx",  # 导出的文件
        input_names=["input"],  # 输入tensor的名称
        output_names=["output"],  # 输出tensor的名称
        opset_version=11)  # onnx算子集版本号

# 读取onnx模型
onnx_model=onnx.load("best_model2.onnx")
# # 检查模型格式是否正确
onnx.checker.check_model(onnx_model)
print("onnx模型格式正确!")
四、onnx模型在本地部署并预测

下面的代码是测试一张图片，而如果想调用自己电脑的摄像头去预测，可以直接使用第二段代码，第二段代码是如何打开摄像头的，然后按照第一段代码的图片处理方式和显示方式即可做到实时的预测了！

第一段代码如下：

import onnxruntime as ort
import torch
import cv2
from PIL import Image
from torchvision import transforms

# 选择GPU还是CPU
provider = ort.get_available_providers()[1 if ort.get_device() == 'GPU' else 0]
# 载入onnx模型
onnx_model = ort.InferenceSession("best_model1.onnx")

# 构造随机输入，获取结果看是否正确
x = torch.randn(1, 3, 224, 224).numpy()

#  输出是一个列表，里面只有一个元素，元素类型是2维数组元素，存着每个分类的概率
# [array([[-0.09820878, -1.3899375 ,  1.5437527 , -0.5629988 ,  0.10385571]],dtype=float32)]
model_output = onnx_model.run(["output"], {"input": x})[0]

# 图像预处理,图像预处理方式要和模型训练时候的预处理方式保持一致
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()
# print(input_tensor.shape) # (1, 3, 224, 224)
# {'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4}

# 对预测的值进行softmax，得到置信度
pred = onnx_model.run(["output"], {"input": input_tensor})[0]
# print(pred.shape) # (1, 5)
pred_softmax = torch.softmax(torch.tensor(pred), dim=1)
# 置信度前三个结果
n = 3
values, indices = torch.topk(pred_softmax, n)  # 返回的是结果，以及这个结果所在原来tensor中的索引位置
# 预测类别
index_n = indices.tolist()[0]
values_n = values.tolist()[0]

src = cv2.imread("test.png")
labels_dict = {'bulbasaur': 0, 'charmander': 1, 'mewtwo': 2, 'pikachu': 3, 'squirtle': 4}
revers_dict = {v: k for k, v in labels_dict.items()}
position = 30
for i in index_n:
    print(revers_dict[i], ":", round(values_n[index_n.index(i)] * 100, 5), "%")
    # 在原图的左上角显示前三个类别的概率
    cv2.putText(src,
                revers_dict[i] + " : " + str(round(values_n[index_n.index(i)] * 100, 5)) + "%",
                (10, position),
                cv2.FONT_HERSHEY_SIMPLEX,
                1, (0, 255, 0), 2)
    position += 30

cv2.imshow("src", src)
cv2.waitKey()
cv2.destroyAllWindows()

效果图如下：

可以看到识别率还是挺高的

在这里插入图片描述

第二段代码如下：

import cv2

# 创建一个VideoCapture对象，参数0表示第一个摄像头，通常是笔记本的内置摄像头
cap = cv2.VideoCapture(0)

# 检查摄像头是否成功打开
if not cap.isOpened():
    print("无法打开摄像头")
    exit()

# 循环读取摄像头的每一帧
while True:
    # 读取一帧
    ret, frame = cap.read()
    
    # 如果正确读取帧，ret为True
    if not ret:
        print("无法读取帧")
        break
    
    # 显示帧
    cv2.imshow('Frame', frame)
    
    # 按'q'键退出循环
    if cv2.waitKey(1) == ord('q'):
        break

# 释放摄像头资源
cap.release()
# 关闭所有OpenCV窗口
cv2.destroyAllWindows()

以上完结撒花！！！

在这里插入图片描述

# 分类 # pytorch # 目标检测 # 人工智能 # cnn # python
Logo
松山湖开发者村综合服务平台

助力广东及东莞地区开发者，代码托管、在线学习与竞赛、技术交流与分享、资源共享、职业发展，成为松山湖开发者首选的工作与学习平台

更多推荐

cover

基于YOLOv8深度学习的智慧农业柑橘病害识别检测系统(PyQt5界面+数据集+训练代码)

avatar 松山湖开发者村综合服务平台
cover

【Kafka】与【Hadoop】的集成应用案例深度解析

avatar 松山湖开发者村综合服务平台
cover

SpringBoot 整合 RabbitMQ

avatar 松山湖开发者村综合服务平台