CIFAR-10 Classification

import torch
import torchvision
import torchvision.transforms as transforms

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import os, sys
import random, math, time
import shutil
import numpy as np
import matplotlib.pyplot as plt

LEARNING_RATE = 1e-3
MOMENTUM = 0.9
EPOCHS = 10
BATCH_SIZE = 4
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

def convert_to_imshow_format(image):
    # first convert back to [0,1] range from [-1,1] range
    image = image / 2 + 0.5
    image = image.numpy()
    # convert from CHW to HWC
    # from 3x32x32 to 32x32x3
    return image.transpose(1,2,0)

def train(model, iterator, optimizer, criterion):
    model.train()
    epoch_loss = 0
    correct = 0
    total = 0
    for idx, (inputs, targets) in enumerate(iterator):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)

        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        epoch_loss += loss.item()
        _, predicted = outputs.max(1)                  # probs, index = torch.max(outputs.data, dim=1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()  # correct += (predicted == labels).sum().item()

    return epoch_loss/len(iterator), 100.*correct/total

def evaluate(model, iterator, criterion):
    model.eval()
    epoch_loss = 0
    correct = 0
    total = 0
    with torch.no_grad():
        for idx, (inputs, targets) in enumerate(iterator):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)

            loss = criterion(outputs, targets)

            epoch_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

    return epoch_loss/len(iterator), 100.*correct/total

def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs

class Net(nn.Module):
  
  def __init__(self):
    super().__init__()
    self.conv1 = nn.Conv2d(3, 16, 3, 1, padding=1)
    self.conv2 = nn.Conv2d(16, 32, 3, 1, padding=1)
    self.conv3 = nn.Conv2d(32, 64, 3, 1, padding=1)
    self.fc1 = nn.Linear(4*4*64, 500)
    self.dropout1 = nn.Dropout(0.5)
    self.fc2 = nn.Linear(500, 10)
    
  def forward(self, x):
    x = F.relu(self.conv1(x))
    x = F.max_pool2d(x, 2, 2)
    x = F.relu(self.conv2(x))
    x = F.max_pool2d(x, 2, 2)
    x = F.relu(self.conv3(x))
    x = F.max_pool2d(x, 2, 2)
    x = x.view(-1, 4*4*64)
    x = F.relu(self.fc1(x))
    x = self.dropout1(x)
    x = self.fc2(x)
    return x

transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

transform_train = transforms.Compose([transforms.Resize((32, 32)),
                                      transforms.RandomHorizontalFlip(),
                                      transforms.RandomRotation(10),
                                      transforms.RandomAffine(0, shear=10, scale=(0.8, 1.2)),
                                      transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
                                      transforms.ToTensor(),
                                      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                     ])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
loaders = {
    'train' : torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE, shuffle=True),
    'test'  : torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE, shuffle=False),
}

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

total_train_batch = len(loaders['train'])
total_test_batch = len(loaders['test'])
print('train 총 배치의 수 : {}'.format(total_train_batch))
print('test 총 배치의 수 : {}'.format(total_test_batch))

model = Net().to(device)
print(model)

criterion = nn.CrossEntropyLoss().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

SAVE_PATH = './model_{}_epoch_{}.pth'.format('VGG', EPOCHS)
BEST_PATH = './best_model_{}.pth'.format('VGG')
train_loss_history = []
train_acc_history = []
val_loss_history = []
val_acc_history = []
best_valid_loss = float('inf')

for epoch in range(EPOCHS):
    start_time = time.time()

    train_loss, train_acc = train(model, loaders['train'], optimizer, criterion)
    valid_loss, val_acc = evaluate(model, loaders['test'], criterion)

    end_time = time.time()

    epoch_mins, epoch_secs = epoch_time(start_time, end_time)

    train_loss_history.append(train_loss)
    train_acc_history.append(train_acc)
    val_loss_history.append(valid_loss)
    val_acc_history.append(val_acc)

    print('Epoch: {} | Time: {}m {}s'.format(epoch + 1, epoch_mins, epoch_secs))
    print('\tTrain Loss: {:.3f} | Acc: {:.2f}%'.format(train_loss, train_acc))
    print('\t Val. Loss: {:.3f} | Acc: {:.2f}%'.format(valid_loss, val_acc))

    if valid_loss < best_valid_loss:
        print('Val. loss decreased ({:.3f} --> {:.3f}).  Saving model ...'.format(best_valid_loss, valid_loss))
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), BEST_PATH)

print('Finished Training')
torch.save(model.state_dict(), SAVE_PATH)

SAVE_PATH = './model_{}_epoch_{}.pth'.format('VGG', EPOCHS)
BEST_PATH = './best_model_{}.pth'.format('VGG')

test_loss, test_acc = evaluate(model, loaders['test'], criterion)
print('| Test Loss: {:.3f} | Acc: {:.2f}%'.format(test_loss, test_acc))

x_len = np.arange(len(train_loss_history))
plt.plot(x_len, train_loss_history, marker='.', c='blue', label="Train-set Loss")
plt.plot(x_len, val_loss_history, marker='.', c='red', label="Val-set Loss")
plt.legend(loc='upper right')
plt.grid()
plt.xlabel('epoch')
plt.ylabel('loss')
plt.show()

dataiter = iter(loaders['test'])
images, labels = dataiter.next()
fig, axes = plt.subplots(1, len(images), figsize=(12, 2.5))
for idx, image in enumerate(images):
    axes[idx].imshow(convert_to_imshow_format(image))
    axes[idx].set_title(classes[labels[idx]])
    axes[idx].set_xticks([])
    axes[idx].set_yticks([])

outputs = model(images.to(device)).to(device)
print("outputs: ", outputs, outputs.size())
sm = nn.Softmax(dim=1)
sm_outputs = sm(outputs)
print(sm_outputs)

probs, index = torch.max(sm_outputs, dim=1)
for p, i in zip(probs, index):
    print('{} - {:.4f}'.format(classes[i], p))

total_correct = 0
total_images = 0
confusion_matrix = np.zeros([10,10], int)
with torch.no_grad():
    for idx, (images, labels) in enumerate(loaders['test']):
        images, labels = images.to(device), labels.to(device)
        outputs = model(images).to(device)
        _, predicted = torch.max(outputs.data, 1)
        total_images += labels.size(0)
        total_correct += (predicted == labels).sum().item()
        for i, l in enumerate(labels):
            confusion_matrix[l.item(), predicted[i].item()] += 1

model_accuracy = total_correct / total_images * 100
print('Model accuracy on {0} test images: {1:.2f}%'.format(total_images, model_accuracy))

print('{0:10s} - {1}'.format('Category','Accuracy'))
for i, r in enumerate(confusion_matrix):
    print('{0:10s} - {1:.1f}'.format(classes[i], r[i]/np.sum(r)*100))

fig, ax = plt.subplots(1, 1, figsize=(8, 6))
ax.matshow(confusion_matrix, aspect='auto', vmin=0, vmax=1000, cmap=plt.get_cmap('Blues'))
plt.ylabel('Actual Category')
plt.yticks(range(10), classes)
plt.xlabel('Predicted Category')
plt.xticks(range(10), classes)
plt.show()

---

Model accuracy on 10000 test images: 75.05%

 

Category - Accuracy

plane - 83.3

car - 82.8

bird - 62.9

cat - 57.9

deer - 71.8

dog - 64.0

frog - 80.4

horse - 78.9

ship - 88.5

truck - 80.0