pytorch 猫狗数据集识别（一）

一 、猫狗数据集预处理

该项目从数据预处理–网络构建–训练–测试这几个部分讲解。读完该部分代码希望能对你有所帮助。 猫狗数据集下载地址链接：https://pan.baidu.com/s/1OBDj51wp0TMRZK_ve-n3fA
提取码：rift

1、数据预处理中加载的模块

import torch
import os
import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader

2、数据预处理代码

class MyDataSet(Dataset):
    def __init__(self, path):
        """根据路径获得数据集"""
        self.path = path
        
        """读取路径下的所有文件"""
        self.dataset = os.listdir(self.path)
        
        # 均值和方差的求解过程下面会交代
        self.mean = [0.4878, 0.4545, 0.4168]
        self.std = [0.2623, 0.2555, 0.2577]

    def __len__(self):
        """获取数据集的长度"""
        return len(self.dataset)

    def __getitem__(self, index):
        """根据索引获得数据集"""
        # 获得数据的标签
        name = self.dataset[index]

        # 数据集的文件名为i.j.jpeg, 这个标签的第i个位置如果是0表示猫，1则表示狗，第j个位置是图片的个数
        name_list = name.split(".")
        target = int(name_list[0])
        
       # 这里需要注意，tensor(target) 和Tensor([target])的区别，否则在one-hot解码的时候会出现错误
        target = torch.tensor(target)

        """数据预处理"""
        # 打开图片
        img = Image.open(os.path.join(self.path, name))
       
        # 设置图片大小
        # img = np.resize(img, (3, 10, 10))
       
        # 归一化，先将图片转化为一个矩阵，然后除以255
        img = np.array(img) / 255

        # # 去均值
        img = (img - self.mean) / self.std

        # 换轴 H W C 转换为 C H W ，这里需要注意下，其实我们可以不这么处理，在前面设置图片大小的时候设置为3 * 100 * 100 的就可以。
        train_data = torch.tensor(img, dtype=torch.float32).permute(2, 0, 1)
        return train_data, target

3、 数据预处理中均值的求解步骤：

if __name__ == '__main__':
    """注意：求均值和标准差时，首先把上面去均值的那一行代码注释了"""
    dataset = MyDataSet(r'C:\Users\admin\Desktop\cat_dog\cat_dog\img')
    
    # 求均值，将DataLoader进行打包，这里batch_size的数量应该是训练图片的总数
    data_loader = DataLoader(dataset= dataset, batch_size=12000, shuffle=True)
    
    # 将data_loader作为一个迭代器进行运算，这里面包含了输入数据和标签，所以后面取第0个位置的数，也就是取输入数据
    data = next(iter(data_loader))[0]

    # 这里图片N(每批的数量) C（通道数） H（高）  W（宽） 几个维度，求的.0是通道层面的均值，所以dim=(0, 2, 3)
    mean = torch.mean(data, dim=(0, 2, 3))
    std = torch.std(data, dim=(0, 2, 3))

二、网络构建

import torch.nn as nn
class MyNetwork(nn.Module):
    def __init__(self):
        super(MyNetwork, self).__init__()
        
        """
        nn.Sequential:一个有序的容器，神经网络模块将按照在传入构造器的顺序依次被添加到计算图中执行
        
        """
        self.line = nn.Sequential(
            nn.Linear(in_features=3 * 100 * 100, out_features=512),
            nn.ReLU(),
            nn.Linear(in_features=512, out_features=256),
            nn.ReLU(),
            nn.Linear(in_features=256, out_features=128),
            nn.ReLU(),
            nn.Linear(in_features=128, out_features=256),
            nn.ReLU(),
            nn.Linear(in_features=256, out_features=512),
            nn.ReLU(),
            nn.Linear(in_features=512, out_features=256),
            nn.ReLU(),
            nn.Linear(in_features=256, out_features=2),
        )

    def forward(self, parse):
        data = torch.reshape(parse, shape=(-1, 3 * 100 * 100))
        return self.line(data)

三、模型的训练

class Train(object):
    def __init__(self, path):
        self.path = path
        self.test_dataset = MyDataSet(os.path.join(self.path, "test_img"))
        self.train_dataset = MyDataSet(os.path.join(self.path, "train_img"))
        self.criterion = torch.nn.MSELoss()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.net = MyNetwork().to(self.device)
        self.optimize = torch.optim.Adam(self.net.parameters())

    def dataLoader(self, batch):
        train_data_loader = DataLoader(dataset=self.train_dataset, batch_size=batch, shuffle=True)
        test_data_loader = DataLoader(dataset=self.test_dataset, batch_size=batch, shuffle=True)
        return train_data_loader, test_data_loader

    def trainNet(self, batch, epoch):
        train_data_loader, test_data_loader = self.dataLoader(batch)
        losses = []
        accuracy = []
        for i in range(epoch):
            for j, (Input, target) in enumerate(train_data_loader):
                Input = Input.to(self.device)
                out_put = self.net(Input)
                """
                one - hot 编码：
                torch.scatter(self, dim, index, value)
                该代码表示把value按照index根据dim的方向填入self中
                Example:
                    index = torch.tensor([0, 2, 4, 1])
                    src = torch.zeros(4, 5)
                    print(src.shape)
                    out = torch.scatter(src, dim=1, index=index.view(-1, 1), value=1)
                    print(out)
                
                """
                target = torch.zeros(target.size(0), 2).scatter_(1, target.view(-1, 1), 1).to(self.device)
                loss = self.criterion(out_put, target)
                losses.append(loss.item())
                
                """ 梯度清零"""
                self.optimize.zero_grad()
                loss.backward()
                self.optimize.step()
                if j % 5 == 0:
                    print("训练轮次：epoch{}/{}，迭代次数：iteration{}/{}".format(i, epoch, j, len(train_data_loader)))
                    acc = torch.mean((out_put.argmax(1) == target.argmax(1)), dtype=torch.float32)
                    accuracy.append(acc.item())
                    print("训练准确率为：accuracy = %s , loss值为：%s" % (acc.item(), loss.item()))
                    plt.clf()
                    plt.ion()
                    """加载loss曲线"""
                    plt.subplot(2, 1, 1)
                    plt.tight_layout(2)
                    plt.plot(losses, label="loss")
                    plt.legend(loc='best')

                    """加载accuracy曲线"""
                    plt.subplot(2, 1, 2)
                   
                    """设置图像之间的距离"""
                    plt.tight_layout(2)
                    plt.plot(accuracy, label="accuracy")
                    """显示图例"""
                    plt.legend(loc='best')
                    plt.pause(0.2)
                    plt.show()
            torch.save(self.net, "model/net.pth")

四、模型的测试

  # 测试
  class Test(object):
    def __init__(self, path):
        self.path = path
        self.test_dataset = MyDataSet(os.path.join(self.path, "test_img"))
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    def Test(self, batch):
        train_data_loader, test_data_loader = self.dataLoader(batch)
        plt.ion()
        for k, (In, tgt) in enumerate(test_data_loader):
            print(k)
            In = In.to(self.device)
            """加载测试图片"""
            # 加载训练好的网络
            test_Net = torch.load('model/net.pth')
            
            # 这是可视化训练图片的第一步，进行反算（Input * 标准差+均值）*255
            # 这里输入In的维度是（batch, 3, 100, 100 ），std 和mean是矩阵，所以给它转换维度才能与输入做运算
            img = (In.cpu().data.numpy() * np.array(self.train_dataset.std).reshape(1, 3, 1, 1) + np.array(self.train_dataset.mean).reshape(1, 3, 1, 1)) * 255
            img = np.resize(img, (100, 100))

            """显示测试结果"""
            if 0 < k <= batch:
                plt.subplot(3, 2, k)
                plt.tight_layout(1)
                plt.imshow(img)
                predict = test_Net(In)
                # print("predict", torch.argmax((predict[k])).item())
                predict_result = "cat" if torch.argmax(predict[k-1]).item() == 0 else "dog"
                plt.title(predict_result)
                plt.axis('off')
                plt.pause(1)
                plt.show()
                plt.savefig("result.jpg")
            elif k == 0:
                continue
            else:
                print("图片测试完毕，请查看！")
                break

五、执行代码

if __name__ == '__main__':
    path = r"数据集路径"
    # 训练
    t = Train(path)
    
    # 测试
    # t.Test(batch=6)