基于pytorch实现猫狗分类系统

本文参加新星计划人工智能(Pytorch)赛道：https://bbs.csdn.net/topics/613989052

写在前言：本文是一个保姆级的分类教程，旨在让零基础的同学掌握实现一个分类系统的基本要素、通用模板和模块实现。在项目代码中做了非常详细的注释，所以就只在文中阐述分类系统实现的步骤，具体解释可查看代码。

下载猫狗数据集:

链接：https://pan.baidu.com/s/18LxKlQV4HAcGxdqYxRhUtg

提取码：3l6q

将数据划分为train、val、test

其中train文件夹下的图片展示，包括猫和狗的图片，实际上数据集总共包含25000张图片，我是以7：2：1的划分方式将数据集分别写入train、val、test中，当然要注意每个数据集下猫狗的数量要均衡否则训练出的模型可能会偏向于预测为某一类。

--dog_cat_classify--data_process.py--model.py--my_loader.py--predict.py--train.py

--data_process.py：用来生成train.txt和val.txt的

--my_laoder.py:自定义Dataset

--model.py:构建自己的猫狗分类模型

--predict.py：使用训练好的模型预测图像

--train.py:使用训练脚本训练并做验证

将train和val里的图片路径和标签以空格隔开，分别保存到train.txt和val.txt中，以备后续自定义Dataset使用。

当自定义Dataset时，Dataset类的初始化函数里的参数就是train.txt或者val.txt的路径，作用是主要将txt里的内容转换为列表的形式

class dog_dataset(nn.Module):def __init__(self, path):super(dog_dataset, self).__init__()with open(path, 'r') as f:self.data_li = f.readlines()

import os#根据图片存放路径生成每张图片的路径和该图片所属的类别，并使用空格分开
def create_data_txt(path, txt_p):data_li = os.listdir(path)with open(txt_p, 'w') as f:for ele in data_li:# print(ele.split('.')[0])if ele.split('.')[0] == 'dog':f.write('%s %s\n'%(os.path.join(path, ele), str(0)))else:f.write('%s %s\n'%(os.path.join(path, ele), str(1)))#可以分别得到train.txt、val.txt用于后面构建自己的数据集
path = r'G:/data_dog_cat/'
txt_path = './'
li = ['train', 'val']
for ele in li:txt_p = os.path.join(txt_path,ele+'.txt')create_data_txt(os.path.join(path, ele), txt_p)

执行完上面脚本后会生成train.txt,val.txt,内容如下面图片所示：

自定义Dataset

import torch.nn as nn
from torchvision.transforms import Compose, ToTensor, Resize, Normalize,ColorJitter
from PIL import Image#实现自己的数据集
class dog_dataset(nn.Module):def __init__(self, path):super(dog_dataset, self).__init__()#数据集集存放的路径self.path = path#设计数据集送入网络前所需要的经过的变换，可以自行组合变换，或者自定义self.transforms = Compose([ColorJitter(),Resize([224, 224]),ToTensor(),Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])#读取对应的txt文件，得到包含所有图像路径和标签的列表with open(path, 'r') as f:self.data_li = f.readlines()def __getitem__(self, index):#通过索引得到列表中图像的路径img_path = self.data_li[index].split(' ')[0]#读取图像，这里要使用PIL库来读取图像，因为后面的transforms里封装的类实例，默认操作的类型是Image类型img = Image.open(img_path)#对图像进行变换，譬如改变图像尺寸、颜色变换、中心裁剪（ToTensor()和Normalize()是必须的，剩余的操作自己组合就好，ToTensor()操作和Normalize()# 操作不可以颠倒，因为Normalize()操作的数据类型是Tensor，不然报类型错误）img = self.transforms(img)#取出图像对应的标签label = int(self.data_li[index].split(' ')[1])return img, labeldef __len__(self):#返回数据集的大小return len(self.data_li)if __name__ == '__main__':txt_path = 'train.txt'mydata = dog_dataset(txt_path)img, lab = mydata[100]print(img.size())print(lab)print(len(mydata))

__init__：主要定义自定义的变换方式，将txt中的图片路径和标签信息以列表的形式存储起来

__getitem__: 主要通过index得到某一条数据，这里是通过index得到列表中的图像路径和标签，然后通过路径读取图像，并利用__init__方法中定义的transform对其进行变换，最后返回图像和标签。

__len__: 返回train或val数据集的大小

构建分类模型

该模型主要通过nn.Sequential的方式进行构建的，nn.Sequential的好处就是可以将几个算子组成一个

模块，模块里的内容会按照顺序执行。

import torch
import torch.nn as nn#自定义一个模型来实现猫狗分类
class myModel(nn.Module):def __init__(self):#因为自定义类继承自nn.Module,所以下面的这句表示初始化父类的初始化方法，这句是必须要有的super(myModel, self).__init__()#本文任务是二分类，最后的输出需要通过激活函数，得到一个概率分布，以0.5为界限，大于等于0.5预测为猫，否则为狗。当然也可以使用softmax()self.sigmoid = nn.Sigmoid()#第一个卷积快包括卷积、激活函数（用于增加模型的非线性也称复杂性，可以增加模型的拟合能力）、归一化层（规范化数据分布的，可以让模型更快的收敛）self.conv1_1 = nn.Sequential(nn.Conv2d(3, 32, (3,3), 1, 1),nn.ReLU(),nn.BatchNorm2d(32))#同上，不同的是卷积步长为2,这个作用可以使图像减小self.conv1_2 = nn.Sequential(nn.Conv2d(32, 64, (3, 3), 2, 1),nn.ReLU(),nn.BatchNorm2d(64))#重复两次，除了参数不同，结构都一样，这样其实可以将模块封装起来进行复用，这里为了更清晰表示，没有进行封装self.conv2_1 = nn.Sequential(nn.Conv2d(64, 64, (3, 3), 1, 1),nn.ReLU(),nn.BatchNorm2d(64))self.conv2_2 = nn.Sequential(nn.Conv2d(64, 128, (3, 3), 2, 1),nn.ReLU(),nn.BatchNorm2d(128))self.conv3_1 = nn.Sequential(nn.Conv2d(128, 128, (3, 3), 1, 1),nn.ReLU(),nn.BatchNorm2d(128))self.conv3_2 = nn.Sequential(nn.Conv2d(128, 256, (3, 3), 2, 1),nn.ReLU(),nn.BatchNorm2d(256))#需要注意的是在卷积和线性层之间要衔接好输入通道数和输出通道数，否则会报错，# 这里的输入通道是将[batch_szie, channel, height, width] --> [batch_size, channel*height*width]self.linear_1 = nn.Linear(28*28*256, 128)self.linear_2 = nn.Linear(128, 1)def forward(self, x):in_size = x.size(0)#这里为了方便阐述张量形状的变化，假定batch_szie为2#[2,3,224,224] --> [2,32,112,112]x = self.conv1_1(x)#[2,32,224,224] --> [2,64,112,112]x = self.conv1_2(x)#[2,64,112,112] --> [2,64,112,112]x = self.conv2_1(x)#[2,64,112,112] --> [2,128,56,56]x = self.conv2_2(x)#[2,64,112,112] --> [2,128,56,56]x = self.conv3_1(x)#[2, 64, 112, 112] --> [2, 256, 28, 28]x = self.conv3_2(x)x = x.view(in_size,-1)x = self.linear_1(x)out = self.linear_2(x)out = self.sigmoid(out)return outif __name__ == '__main__':x = torch.rand([2, 3, 224, 224])model = myModel()print(model)print(model(x).size())print(model(x))

__init__: 初始化算子

__forward__: 真正的算子调用是在这个函数里进行的，通过输入x，得到各个算子的计算最后得到概率分布的形式

最后打印一下可以看到模型的结构：

myModel((sigmoid): Sigmoid()(conv1_1): Sequential((0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(1): ReLU()(2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))(conv1_2): Sequential((0): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))(1): ReLU()(2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))(conv2_1): Sequential((0): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(1): ReLU()(2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))(conv2_2): Sequential((0): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))(1): ReLU()(2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))(conv3_1): Sequential((0): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(1): ReLU()(2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))(conv3_2): Sequential((0): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))(1): ReLU()(2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))(linear_1): Linear(in_features=200704, out_features=128, bias=True)(linear_2): Linear(in_features=128, out_features=1, bias=True)
)

搭建训练脚本

import torch
from torch.utils.data import DataLoader
from .my_loader import dog_dataset
from .model import myModel
from torch import nn
from torch.optim.lr_scheduler import StepLR# 定义训练过程
def train(model, device, train_loader, optimizer, epoch, lr_scheduler, criterion):#定义模型的初始化模式，train和val时，batchnorm和dropout的使用方式不同model.train()#迭代dataloader,为什么datlaoder可以被遍历(因为它是可迭代对象)for batch_idx, (data, target) in enumerate(train_loader):data, target = data.to(device), target.to(device).float().unsqueeze(1)#梯度清零optimizer.zero_grad()#前向传播output = model(data)#计算lossloss = criterion(output, target)#反向传播loss.backward()#参数更新optimizer.step()#学习率更新lr_scheduler.step()#每迭代10次，打印一次lossif (batch_idx + 1) % 10 == 0:print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),(batch_idx + 1) / len(train_loader), loss.item()))# 定义测试过程def val(model, device, test_loader, criterion):#定义训练模式model.eval()test_loss = 0correct = 0#上下文管理器，在此范围内不计算梯度with torch.no_grad():for data, target in test_loader:data, target = data.to(device), target.to(device).float().unsqueeze(1)output = model(data)# print(output)test_loss += criterion(output, target, reduction='mean').item()pred = torch.tensor([[1] if num[0] >= 0.5 else [0] for num in output]).to(device)correct += pred.eq(target.long()).sum().item()#打印一次准确率print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(test_loss, correct, len(test_loader.dataset),correct / len(test_loader.dataset)))def main():# 如果有cuda则使用cuda加速训练模型，否则就是用gpudevice = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')#定义训练轮数epoches = 100# 实例化训练数据集train_dataset = dog_dataset('C:/Users/86181/Desktop/tset_demo\dog_cat_classify/train.txt')# 实例化验证数据集val_dataset = dog_dataset('C:/Users/86181/Desktop/tset_demo\dog_cat_classify/val.txt')# 实例化训练的dataLoadertrain_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)# 实例化验证的dataLoaderval_loader = DataLoader(val_dataset, batch_size=2, shuffle=True)# 实例化模型model = myModel()# 实例化损失criterion = nn.BCELoss()# 实例化一个优化器,初始化学习率为1e-3optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)# 实例化一个自适应学习率策略lr_scheduler = StepLR(optimizer, step_size=10, gamma=0.1)for epoch in range(1, epoches+1):train(model, device, train_loader, optimizer, epoch, lr_scheduler, criterion)val(model, device, val_loader, criterion)#保存模型torch.save(model, 'G:/dog_cat_calssify/model.pth')

训练情况如下图所示：

搭建预测脚本

import torch
from PIL import Image
from torchvision import transformsdef predict(model_save_path, device, img_path):class_names = ['dog', 'cat']model = torch.load(model_save_path)model.eval()image_PIL = Image.open(img_path)transform_test = transforms.Compose([transforms.Resize(224),transforms.ToTensor(),transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])image_tensor = transform_test(image_PIL)image_tensor = torch.unsqueeze(image_tensor, 0)image_tensor = image_tensor.to(device)out = model(image_tensor)pred = torch.tensor([[1] if num[0] >= 0.5 else [0] for num in out]).to(device)return class_names[pred]