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

本文主要是介绍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 picklewith 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 timeclass 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 = Noneself.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 = dataelse: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, labeldef __len__(self):return len(self.labels)def _unpickle(file):import picklewith 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 nnclass 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学习率变为 $l r * g a m m a$ 。

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 DataLoadertrainloader = 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 = 0t1 = time.time()for i, (data, label) in enumerate(trainloader):data, label = data.to(device), label.to(device)# defaultoptimizer.zero_grad()out = model(data)loss = criterion(out, label)loss.backward()optimizer.step()# print infodata_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_lengthprint('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 = 0with torch.no_grad():for i, (data, label) in enumerate(testloader):data, label = data.to(device), label.to(device)# defaultout = 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_lengthprint('TestLoss:{:.4f}, Acc@1:{:.4f}'.format(running_loss, acc), end=' ')# save modelglobal best_accif acc > best_acc:best_acc = accstate = {'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))