利用PyTorch自定义数据集实现猫狗分类

看了许多关于PyTorch的入门文章，大抵是从torchvision.datasets中自带的数据集进行训练，导致很难把PyTorch运用于自己的数据集上，真正地灵活运用PyTorch。

这里我采用从Kaggle上下载的猫狗数据集，利用自定义数据集训练自己的二分类神经网络。

解压后，一个文件里面有12500张图，猫狗各一半，文件名类似于这样：cat.0.jpg、dog.12499.jpg

因为只是练手，所以不用这么大的，仅仅采用子数据集。

利用Python的os库，将数据集进行拆分。分为train与test两个文件架，每个里面都有cats和dogs两个文档。train里面每种动物有1000张图，test里面每种动物有500张图。图片大概是这个样子（大小不一）：

接下里开始编码：

# 导入库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

# 设置超参数
BATCH_SIZE = 50
EPOCHS = 30
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 数据预处理
transform = transforms.Compose([
    transforms.RandomResizedCrop(150),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])

# 读取数据
root = 'Cats_Dogs'
dataset_train = datasets.ImageFolder(root + '\\train', transform)
dataset_test = datasets.ImageFolder(root + '\\test', transform)

# 导入数据
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=True)

# 定义网络
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.max_pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.max_pool2 = nn.MaxPool2d(2)
        self.conv3 = nn.Conv2d(64, 128, 3)
        self.max_pool3 = nn.MaxPool2d(2)
        self.conv4 = nn.Conv2d(128, 128, 3)
        self.max_pool4 = nn.MaxPool2d(2)
        self.fc1 = nn.Linear(6272, 512)
        self.fc2 = nn.Linear(512, 1)
        
    def forward(self, x):
        in_size = x.size(0)
        x = self.conv1(x)
        x = F.relu(x)
        x = self.max_pool1(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.max_pool2(x)
        x = self.conv3(x)
        x = F.relu(x)
        x = self.max_pool3(x)
        x = self.conv4(x)
        x = F.relu(x)
        x = self.max_pool4(x)
        # 展开
        x = x.view(in_size, -1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = torch.sigmoid(x)
        return x

# 实例化模型并且移动到GPU
model = ConvNet().to(DEVICE)
# 选择简单暴力的Adam优化器，学习率调低
optimizer = optim.Adam(model.parameters(), lr=1e-4)

# 定义训练过程
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        # 由于评论区小伙伴的提醒，我现在发现这里更好的做法是利用unsqueeze()增加一个维度而不是reshape，否则要是提取的不是50的倍数就会报错
        data, target = data.to(device), target.to(device).float().reshape(50, 1)
        optimizer.zero_grad()
        output = model(data)
        # print(output)
        loss = F.binary_cross_entropy(output, target)
        loss.backward()
        optimizer.step()
        if(batch_idx+1)%10 == 0: 
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, (batch_idx+1) * len(data), len(train_loader.dataset),
                100. * (batch_idx+1) / len(train_loader), loss.item()))

# 定义测试过程
def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device).float().reshape(50, 1)
            output = model(data)
            # print(output)
            test_loss += F.binary_cross_entropy(output, target, reduction='sum').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),
        100. * correct / len(test_loader.dataset)))

# 训练
for epoch in range(1, EPOCHS + 1):
    train(model, DEVICE, train_loader, optimizer, epoch)
    test(model, DEVICE, test_loader)

在2000张图片上，我们利用小的深度神经网络，训练出了一个正确率为72%的分类器。虽然结果不太理想，但是在没有进行任何防止过拟合操作的情况下，还算是过得去的成绩。如果添加Dropout或者正则化、数据增强的话，相信结果会有不错的提升。而我们使用的数据才是整个数据集的很小一部分而已。