【Kaggle猫狗大战】Pytorch 自定义神经网络实现简单的二分类问题

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 如下
image.png
设计好卷积网络后,在卷积网络后面跟了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_looptest_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. 保存模型
阅读(130)
评论(0)
updated@2021-05-07
评论区
目录