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进行深度学习项目大体就是这么个流程:
- 准备数据,自定义Dataset
- 设计神经网络
- 选取合适的损失函数和优化函数
- 训练模型
- 保存模型