【小白学习PyTorch教程】八、使用图像数据增强手段，提升CIFAR-10 数据集精确度
「@Author：Runsen」
上次基于CIFAR-10 数据集，使用PyTorch构建图像分类模型的精确度是60%，对于如何提升精确度，方法就是常见的transforms图像数据增强手段。
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import torchvision.utils as vutils
import numpy as np
import os
import warnings
from matplotlib import pyplot as plt
warnings.filterwarnings('ignore')`
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
加载数据集
# number of images in one forward and backward pass
batch_size = 128
# number of subprocesses used for data loading
# Normally do not use it if your os is windows
num_workers = 2
train_dataset = datasets.CIFAR10('./data/CIFAR10/',
train = True,
download = True,
transform = transform_train)
train_loader = DataLoader(train_dataset,
batch_size = batch_size,
shuffle = True,
num_workers = num_workers)
val_dataset = datasets.CIFAR10('./data/CIFAR10',
train = True,
transform = transform_test)
val_loader = DataLoader(val_dataset,
batch_size = batch_size,
shuffle = False,
num_workers = num_workers)
test_dataset = datasets.CIFAR10('./data/CIFAR10',
train = False,
transform = transform_test)
test_loader = DataLoader(test_dataset,
batch_size = batch_size,
shuffle = False,
num_workers = num_workers)
# declare classes in CIFAR10
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
之前的transform ’只是进行了缩放和归一，在这里添加RandomCrop和RandomHorizontalFlip
# define a transform to normalize the data
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), # converting images to tensor
transforms.Normalize(mean = (0.5, 0.5, 0.5), std = (0.5, 0.5, 0.5))
# if the image dataset is black and white image, there can be just one number.
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean = (0.5, 0.5, 0.5), std = (0.5, 0.5, 0.5))
])
可视化具体的图像
# function that will be used for visualizing the data
def imshow(img):
img = img / 2 + 0.5 # unnormalize
plt.imshow(np.transpose(img, (1, 2, 0))) # convert from Tensor image
# obtain one batch of imges from train dataset
dataiter = iter(train_loader)
images, labels = dataiter.next()
images = images.numpy() # convert images to numpy for display
# plot the images in one batch with the corresponding labels
fig = plt.figure(figsize = (25, 4))
# display images
for idx in np.arange(10):
ax = fig.add_subplot(1, 10, idx+1, xticks=[], yticks=[])
imshow(images[idx])
ax.set_title(classes[labels[idx]])
建立常见的CNN模型
# define the CNN architecture
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.main = nn.Sequential(
# 3x32x32
nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 3, padding = 1), # 3x32x32 (O = (N+2P-F/S)+1)
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size = 2, stride = 2), # 32x16x16
nn.BatchNorm2d(32),
nn.Conv2d(32, 64, kernel_size = 3, padding = 1), # 32x16x16
nn.ReLU(inplace=True),
nn.MaxPool2d(2, 2), # 64x8x8
nn.BatchNorm2d(64),
nn.Conv2d(64, 128, 3, padding = 1), # 64x8x8
nn.ReLU(inplace=True),
nn.MaxPool2d(2, 2), # 128x4x4
nn.BatchNorm2d(128),
)
self.fc = nn.Sequential(
nn.Linear(128*4*4, 1024),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(1024, 256),
nn.ReLU(inplace=True),
nn.Dropout(0.5),
nn.Linear(256, 10)
)
def forward(self, x):
# Conv and Poolilng layers
x = self.main(x)
# Flatten before Fully Connected layers
x = x.view(-1, 128*4*4)
# Fully Connected Layer
x = self.fc(x)
return x
cnn = CNN().to(device)
cnn
torch.nn.CrossEntropyLoss对输出概率介于0和1之间的分类模型进行分类。
训练模型
# 超参数：Hyper Parameters
learning_rate = 0.001
train_losses = []
val_losses = []
# Loss function and Optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(cnn.parameters(), lr = learning_rate)
# define train function that trains the model using a CIFAR10 dataset
def train(model, epoch, num_epochs):
model.train()
total_batch = len(train_dataset) // batch_size
for i, (images, labels) in enumerate(train_loader):
X = images.to(device)
Y = labels.to(device)
### forward pass and loss calculation
# forward pass
pred = model(X)
#c alculation of loss value
cost = criterion(pred, Y)
### backward pass and optimization
# gradient initialization
optimizer.zero_grad()
# backward pass
cost.backward()
# parameter update
optimizer.step()
# training stats
if (i+1) % 100 == 0:
print('Train, Epoch [%d/%d], lter [%d/%d], Loss: %.4f'
% (epoch+1, num_epochs, i+1, total_batch, np.average(train_losses)))
train_losses.append(cost.item())ｎ
# def the validation function that validates the model using CIFAR10 dataset
def validation(model, epoch, num_epochs):
model.eval()
total_batch = len(val_dataset) // batch_size
for i, (images, labels) in enumerate(val_loader):
X = images.to(device)
Y = labels.to(device)
with torch.no_grad():
pred = model(X)
cost = criterion(pred, Y)
if (i+1) % 100 == 0:
print("Validation, Epoch [%d/%d], lter [%d/%d], Loss: %.4f"
% (epoch+1, num_epochs, i+1, total_batch, np.average(val_losses)))
val_losses.append(cost.item())
def plot_losses(train_losses, val_losses):
plt.figure(figsize=(5, 5))
plt.plot(train_losses, label='Train', alpha=0.5)
plt.plot(val_losses, label='Validation', alpha=0.5)
plt.xlabel('Epochs')
plt.ylabel('Losses')
plt.legend()
plt.grid(b=True)
plt.title('CIFAR 10 Train/Val Losses Over Epoch')
plt.show()
num_epochs = 20
for epoch in range(num_epochs):
train(cnn, epoch, num_epochs)
validation(cnn, epoch, num_epochs)
torch.save(cnn.state_dict(), './data/Tutorial_3_CNN_Epoch_{}.pkl'.format(epoch+1))
plot_losses(train_losses, val_losses)
测试模型
def test(model):
# declare that the model is about to evaluate
model.eval()
correct = 0
total = 0
with torch.no_grad():
for images, labels in test_dataset:
images = images.unsqueeze(0).to(device)
# forward pass
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += 1
correct += (predicted == labels).sum().item()
print("Accuracy of Test Images: %f %%" % (100 * float(correct) / total))
经过图像数据增强。模型从60提升到了84。
测试模型在哪些类上表现良好，
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
for data in test_loader:
images, labels = data
images = images.to(device)
labels = labels.to(device)
outputs = cnn(images)
_, predicted = torch.max(outputs, 1)
c = (predicted == labels).squeeze()
for i in range(4):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1
for i in range(10):
print('Accuracy of %5s : %2d %%' % (
classes[i], 100 * class_correct[i] / class_total[i]))
评论