1. 项目介绍

猫狗大战是Kaggle很早的机器学习入门赛了。项目提供25000张图片（猫狗分别12500张）作为训练数据，是不错的练手项目，可以了解一下使用 Python 和Pytorch 准备数据，设计神经网络模型，训练模型， 导出模型和模型推理的整个过程。

2. 准备数据

2.1. 下载猫狗大战的数据集

Kaggle官方地址：点我下载（可能要科学上网）
热心网友提供的百度云分享链接：https://pan.baidu.com/s/1kqfkr2X7mMkuFXb6C3KgTg 密码：xzyh

2.2. 划分数据集和验证集

我们通过随机抽样，将10%的数据样本作为验证集

import os
import glob
import random
import shutil

def mkdirs_if_missing(path):
    if not os.path.exists(path):
        os.makedirs(path)

# 划分数据集
def split_data(rate=0.1):
    
    src = './data\\train'
    
    # 1.检查并创建文件夹
    val_src ='./data\\val'
    mkdirs_if_missing(val_src)
    
    # 2.抽样
    img_pths = glob.glob(src + os.sep + '**.jpg')
    random.shuffle(img_pths)     # 打乱顺序
    val_count = int(len(img_pths) * rate)
    val_pths = random.sample(img_pths, val_count)    # 抽样
    # 3.移动抽样的图片到验证集
    for pth in val_pths:
        name = pth.split(os.sep)[-1]
        shutil.move(pth, os.path.join(val_src, name))
    print(f'验证集样本数量为：{val_count}')

2.3. 自定义我们的 Dataset

根据官网文档介绍 Dataset 需要要实现三个方法

• init(): 初始化一个Dataset
• len()： 返回样本数
• getitem()： 给定一个参数idx， 返回该位置的样本和对应的label
  由于猫狗大战的训练样本 label 是标注在图片名称上的。所以我们通过图片名字来判断label
  Dog 为 1, Cat 为 0

import torch
from torch.utils.data import Dataset
import cv2 as cv
import glob
import os

class DogCatDataset(Dataset):
    def __init__(self, img_dir):
        self.img_dir = img_dir
        img_pths = glob.glob(img_dir  + os.sep + '**.jpg')
        assert img_pths, 'no jpg file in ' + img_dir
        self.img_pths = img_pths
    
    def __len__(self):
        return len(self.img_pths)
    
    def __getitem__(self, idx):
        img_pth = self.img_pths[idx]
        img_name = img_pth.split(os.sep)[-1]
        
        label = 1 if img_name.split('.')[0] == 'dog' else 0
        label = torch.tensor(label)
        image = cv.imread(img_pth)
        image = cv.resize(image, (360, 360), cv.INTER_LINEAR)    # resize
        image = image / 255.0    # 归一化
        image = torch.from_numpy(image)    # 转为Tensor
        image = image.permute(2, 0, 1).to(torch.float32)
        return image, label

3. 设计神经网络

在 Dataset 中我们把图片通过opencv的resize()方法统一成 360x360 大小。所以我们的神经网络输入为 360x360@3， 3是通道数
首先是设计卷积网络提取图片特征， 我随便设计了一个卷积网络，最终输出 4x4@384 如下

设计好卷积网络后，在卷积网络后面跟了4个全连接层，代码如下

from torch import nn
import torch.nn.functional as F
class DogCatNet(nn.Module):
    def __init__(self):
        super(DogCatNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4)
        self.pool = nn.MaxPool2d(kernel_size=3, stride=2)
        self.conv2 = nn.Conv2d(96, 256, kernel_size=3, stride=1)
        self.conv3 =  nn.Conv2d(256, 384, kernel_size=3, stride=2)
        self.fc1 = nn.Linear(384 * 4 * 4, 256)
        self.fc2 = nn.Linear(256, 64)
        self.fc3 = nn.Linear(64, 32)
        self.fc4 = nn.Linear(32, 2)
        
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 384 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        
        return x

4. 训练数据

首先定义两个方法train_loop和test_loop一个用于训练更新数据，一个用于阶段性验证模型的预测效果

import time

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    t0 = time.time()
    for batch, (X, y) in enumerate(dataloader):
        if torch.cuda.is_available():
            X = X.cuda()
            y = y.cuda()
        # Compute prediction and loss
        pred = model(X)
        loss = loss_fn(pred, y)
        

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (batch + 1) % 10 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}] ({time.time()-t0:.3f}s)")
            t0 = time.time()

def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            # 使用gpu加速
            if torch.cuda.is_available():
                X = X.cuda()
                y = y.cuda()
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= size
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

然后编写代码进行训练，需要注意的是在windows上DataLoader的num_workers不设置为0的化会出错。

# 1. 使用DataLoader和自定义的Dataset加载数据
trainset = DogCatDataset('./data\\train')
testset = DogCatDataset('./data\\val')
train_loader = DataLoader(trainset, batch_size=96, num_workers=0, shuffle=True)
test_loader = DataLoader(testset, batch_size=96, num_workers=0, shuffle=True)

# 2. 创建模型，确定损失函数和优化函数
model = DogCatNet()
loss_fn = nn.CrossEntropyLoss()
# GPU 加速
if torch.cuda.is_available():
    model = model.cuda()
    loss_fn = loss_fn.cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 3. 训练模型
epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_loader, model, loss_fn, optimizer)
    test_loop(test_loader, model, loss_fn)
print("Done!")

这里我通过10个epochs的训练， 在验证集的准确率上达到了 85%

5. 保存模型

torch.save(model, './model_final.pt')

6. 小结

使用Pytorch进行深度学习项目大体就是这么个流程：

1. 准备数据，自定义Dataset
2. 设计神经网络
3. 选取合适的损失函数和优化函数
4. 训练模型
5. 保存模型