# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# STAT 453: Deep Learning (Spring 2021)  
# Instructor: <NAME> (<EMAIL>)  
#
# Course website: http://pages.stat.wisc.edu/~sraschka/teaching/stat453-ss2021/  
# GitHub repository: https://github.com/rasbt/stat453-deep-learning-ss21
#
# ---

# %load_ext watermark
# %watermark -a '<NAME>' -v -p torch

# # All-Convolutional Net on Cifar-10

# ## Imports

import torch
import torchvision
import numpy as np
import matplotlib.pyplot as plt

# From local helper files
from helper_evaluation import set_all_seeds, set_deterministic, compute_confusion_matrix
from helper_train import train_model
from helper_plotting import plot_training_loss, plot_accuracy, show_examples, plot_confusion_matrix
from helper_dataset import get_dataloaders_cifar10, UnNormalize

# ## Settings and Dataset

# +
##########################
### SETTINGS
##########################

RANDOM_SEED = 123
BATCH_SIZE = 256
NUM_EPOCHS = 50
DEVICE = torch.device('cuda:2' if torch.cuda.is_available() else 'cpu')
# -

set_all_seeds(RANDOM_SEED)
#set_deterministic()

# +
##########################
### CIFAR-10 DATASET
##########################

train_transforms = torchvision.transforms.Compose([
    torchvision.transforms.Resize((70, 70)),
    torchvision.transforms.RandomCrop((64, 64)),
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                      ])

test_transforms = torchvision.transforms.Compose([
    torchvision.transforms.Resize((70, 70)),        
    torchvision.transforms.CenterCrop((64, 64)),            
    torchvision.transforms.ToTensor(),                
    torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])




train_loader, valid_loader, test_loader = get_dataloaders_cifar10(
    batch_size=BATCH_SIZE,
    validation_fraction=0.1,
    train_transforms=train_transforms,
    test_transforms=test_transforms,
    num_workers=2)

# Checking the dataset
for images, labels in train_loader:  
    print('Image batch dimensions:', images.shape)
    print('Image label dimensions:', labels.shape)
    print('Class labels of 10 examples:', labels[:10])
    break
# -

# ## Model

# +
##########################
### MODEL
##########################


class AllConvNet(torch.nn.Module):

    def __init__(self, num_classes):
        super().__init__()
        
        self.net = torch.nn.Sequential(
                      torch.nn.Conv2d(in_channels=3,
                                      out_channels=16,
                                      kernel_size=(3, 3),
                                      stride=(1, 1),
                                      padding=1,
                                      bias=False),
                      torch.nn.BatchNorm2d(16),
                      torch.nn.ReLU(inplace=True),
                      torch.nn.Conv2d(in_channels=16,
                                      out_channels=16,
                                      kernel_size=(3, 3),
                                      stride=(2, 2),
                                      padding=1,
                                      bias=False),   
                      torch.nn.BatchNorm2d(16),
                      torch.nn.ReLU(inplace=True),
                      torch.nn.Conv2d(in_channels=16,
                                      out_channels=32,
                                      kernel_size=(3, 3),
                                      stride=(1, 1),
                                      padding=1,
                                      bias=False),        
                      torch.nn.BatchNorm2d(32),
                      torch.nn.ReLU(inplace=True),                        
                      torch.nn.Conv2d(in_channels=32,
                                      out_channels=32,
                                      kernel_size=(3, 3),
                                      stride=(2, 2),
                                      padding=1,
                                      bias=False),      
                      torch.nn.BatchNorm2d(32),
                      torch.nn.ReLU(inplace=True),       
                      torch.nn.Conv2d(in_channels=32,
                                      out_channels=64,
                                      kernel_size=(3, 3),
                                      stride=(1, 1),
                                      padding=1,
                                      bias=False),   
                      torch.nn.BatchNorm2d(64),
                      torch.nn.ReLU(inplace=True),
                      torch.nn.Conv2d(in_channels=64,
                                      out_channels=64,
                                      kernel_size=(3, 3),
                                      stride=(2, 2),
                                      padding=1,
                                      bias=False),     
                      torch.nn.BatchNorm2d(64),
                      torch.nn.ReLU(inplace=True),
                      torch.nn.Conv2d(in_channels=64,
                                      out_channels=num_classes,
                                      kernel_size=(3, 3),
                                      stride=(1, 1),
                                      padding=1,
                                      bias=False),    
                      torch.nn.BatchNorm2d(10),
                      torch.nn.ReLU(inplace=True),
                      torch.nn.AdaptiveAvgPool2d(1),
                      torch.nn.Flatten()
    )


    def forward(self, x):
        
        x = self.net(x)
        #probas = torch.softmax(x, dim=1)
        return x

# +
model = AllConvNet(num_classes=10)

model = model.to(DEVICE)

optimizer = torch.optim.SGD(model.parameters(), momentum=0.9, lr=0.1)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
                                                       factor=0.1,
                                                       mode='max',
                                                       verbose=True)

minibatch_loss_list, train_acc_list, valid_acc_list = train_model(
    model=model,
    num_epochs=NUM_EPOCHS,
    train_loader=train_loader,
    valid_loader=valid_loader,
    test_loader=test_loader,
    optimizer=optimizer,
    device=DEVICE,
    scheduler=scheduler,
    scheduler_on='valid_acc',
    logging_interval=100)

plot_training_loss(minibatch_loss_list=minibatch_loss_list,
                   num_epochs=NUM_EPOCHS,
                   iter_per_epoch=len(train_loader),
                   results_dir=None,
                   averaging_iterations=200)
plt.show()

plot_accuracy(train_acc_list=train_acc_list,
              valid_acc_list=valid_acc_list,
              results_dir=None)
plt.ylim([60, 100])
plt.show()

# +
model.cpu()
unnormalizer = UnNormalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
class_dict = {0: 'airplane',
              1: 'automobile',
              2: 'bird',
              3: 'cat',
              4: 'deer',
              5: 'dog',
              6: 'frog',
              7: 'horse',
              8: 'ship',
              9: 'truck'}

show_examples(model=model, data_loader=test_loader, unnormalizer=unnormalizer, class_dict=class_dict)
# -

mat = compute_confusion_matrix(model=model, data_loader=test_loader, device=torch.device('cpu'))
plot_confusion_matrix(mat, class_names=class_dict.values())
plt.show()

StarcoderdataJupyterScriptsDedupFiltered