PyTorch Demo-1 : CIFAR-10 分类模型

PyTorch >= 1.0

Python 3

1. Data

1.1 CIFAR-10 数据

CIFAR-10 官方下载 ，下载为 cifar-10-python.tar.gz

解压缩文件包含：

‘batches.meta’，‘data_batch_1’，‘data_batch_2’，‘data_batch_3’，‘data_batch_4’，‘data_batch_5’，‘test_batch’

由官方代码读取，其中 batches.meta 中为描述内容 data_batch_* 为训练集，test_batch 为测试集

读取数据：

def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict

数据显示：以测试数据为例，单个数据维度为（3072，）

data = unpickle('test_batch')
"""
data.keys():[b'batch_label', b'labels', b'data', b'filenames']
"""
# 获取单个数据
img = data[b'data'][0].reshape(3, 32, 32).transpose(1, 2, 0)
fname = data[b'filenames'][0]
label = data[b'label'][0]
"""
图片显示: plt.imshow(img)
fname: b'domestic_cat_s_000907.png'
label: 3
"""

1.2 构造Dataset

torchvision 有自带的函数 torchvision.datasets.CIFAR10() 可直接处理CIFAR10数据，此处采用自定义数据集的方式，需要继承 torch.utils.data 下的 dataset.Dataset ，重写 __init__() ，__getitem__() ，__len__() 函数，具体设计根据数据来。如果数据是按照文件夹分好的可以直接使用 torchvision.datasets.ImageFolder() ，详见官网 。

from torch.utils.data import dataset
from torchvision import transforms
import numpy as np
import os
import time

class CIFAR10(dataset.Dataset):
    def __init__(self, mode):
        assert mode in ['train', 'test'], print('mode must be "train" or "test"')
        data_root = './data/cifar-10-batches-py/' # 文件目录
        data_files = {'train': ['data_batch_1', 'data_batch_2', 'data_batch_3', 'data_batch_4', 'data_batch_5'],
                'test': ['test_batch']}
        self.imgs = None
        self.labels = []
        # self.class_names = self._unpickle(os.path.join(data_root, 'batches.meta'))[b'label_names]
        for f in data_files[mode]:
            data_dict = self._unpickle(os.path.join(data_root, f))
            data = data_dict[b'data'].reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
            if self.imgs is None:
                self.imgs = data
            else:
                self.imgs = np.vstack((self.imgs, data))
            self.labels += data_dict[b'labels']
            
        if mode == 'train':
            # 训练集加入随机翻转, 数据增强
            self.trans = transforms.Compose([
                transforms.ToPILImage(),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
        else:
            self.trans = transforms.Compose([
                transforms.ToPILImage(),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
    
    def __getitem__(self, index):
        img = self.imgs[index]
        label = self.labels[index]
        img = self.trans(img)
        
        return img, label
    
    def __len__(self):
        return len(self.labels)
    
    def _unpickle(file):
        import pickle
        with open(file, 'rb') as fo:
            dict = pickle.load(fo, encoding='bytes')
        return dict

2. Model

PyTorch定义模型需要继承nn.Module ，重写 __init__() 和 forward() 函数直接在初始中定义网络需要的结构，前向传播函数定义执行的顺序。此处模型参考 《Binary Classification from Positive Data with Skewed Confidence》 ，在CIFAR10数据集准确率 75% 左右。

import torch
import torch.nn as nn

class CIFAR10_Net(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.feature = nn.Sequential(
            nn.Conv2d(3, 18, kernel_size=5, padding=2, stride=1),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(18, 48, kernel_size=5, padding=2, stride=1),
            nn.ReLU(True),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(48*8*8, 800),
            nn.ReLU(True),
            nn.Linear(800, 400),
            nn.ReLU(True),
            nn.Linear(400, num_classes)
        )

    def forward(self, x):
        out = self.feature(x)
        out = out.view(out.size(0), -1)
        out = self.fc(out)

        return out

3. Train

3.1 初始设置

设置是否使用GPU，以及每个批次的个数和训练次数。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size = 64
total_epoch = 10
best_acc = 0.

3.2 损失函数、优化器等…

损失函数采用交叉熵函数，优化器使用SGD，scheduler 为学习率衰减，设置每 8 个epoch学习率变为 $lr\ast gamma$ 。

import torch.nn as nn
import torch.optim as optim

# model
model = CIFAR10_Net(10).to(device)
# loss
criterion = nn.CrossEntropyLoss()
# optimizer
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=3e-4, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.1)

3.3 DataLoader

DataLoader 用于加载数据集，num_workers 为多线程设置，默认为0，pin_memory 为锁页内存，设置为True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些，如果内存不足则False。

from torch.utils.data import DataLoader

trainloader = DataLoader(dataset=CIFAR10('train'), batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True)
testloader = DataLoader(dataset=CIFAR10('test'), batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True)

3.4 Train & Test & Save model

def train():
    model.train()
    running_loss = 0.
    running_correct = 0.
    data_length = 0
    t1 = time.time()
    for i, (data, label) in enumerate(trainloader):
        data, label = data.to(device), label.to(device)
        # default
        optimizer.zero_grad()
        out = model(data)
        loss = criterion(out, label)
        loss.backward()
        optimizer.step()
        # print info
        data_length += data.size(0)
        running_loss += loss.item() * data.size(0) / data_length
        _, pred = torch.max(out, 1)
        running_correct += pred.eq(label).sum().item()

    acc = running_correct / data_length

    print('Loss:{:.4f}, Acc@1:{:.4f}, time:{:.2f}'.format(running_loss, acc, time.time() - t1), end=' -> ')

def test(epoch):
    model.eval()
    running_loss = 0.
    running_correct = 0.
    data_length = 0
    with torch.no_grad():
        for i, (data, label) in enumerate(testloader):
            data, label = data.to(device), label.to(device)
            # default
            out = model(data)
            loss = criterion(out, label)

            data_length += data.size(0)
            running_loss += loss.item() * data.size(0) / data_length
            _, pred = torch.max(out, 1)
            running_correct += pred.eq(label).sum().item()

    acc = running_correct / data_length
    
    print('TestLoss:{:.4f}, Acc@1:{:.4f}'.format(running_loss, acc), end=' ')
    # save model
    global best_acc
    if acc > best_acc:
        best_acc = acc
        state = {
            'net': model.state_dict(),
            'epoch': epoch,
            'best_acc': best_acc
        }
        torch.save(state, 'ckpt.pth')
        print('*')
    else:
        print()

3.5 Main

t1 = time.time()
for epoch in range(total_epoch):
    print('epoch[{:>3}/{:>3}]'.format(epoch, total_epoch), end=' ')
    train()
    scheduler.step()
    test(epoch)
    
t = time.time() - t1
print('\ntotal time:{}min{:.2f}s, best_acc:{:.4f}'.format(t//60, t%60, best_acc))

训练曲线：

Reference:

[1] TRANSFER LEARNING FOR COMPUTER VISION TUTORIAL .

[2] PYTORCH DOCUMENTATION .

[3] The CIFAR-10 dataset .