【小白学习PyTorch教程】八、使用图像数据增强手段，提升CIFAR-10 数据集精确度

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')

# 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]])

# 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

# 超参数：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))

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]))