Transfer Learning Exercise

# python
from collections import OrderedDict
from datetime import datetime
import os

# pypi
from dotenv import load_dotenv
from torch import nn
from sklearn.model_selection import train_test_split
from torch.utils.data.sampler import SubsetRandomSampler

import matplotlib
import numpy
import seaborn
import torch
import torch.optim as optimize
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as pyplot

# this project
from neurotic.tangles.data_paths import DataPathTwo

load_dotenv()
path = DataPathTwo(folder_key="FLOWERS")
print(path.folder)
for target in path.folder.iterdir():
    print(target)

fig = pyplot.figure(figsize=(12, 10))
pyplot.rc("axes", titlesize=10)
for idx in numpy.arange(20):
    ax = fig.add_subplot(2, 20/2, idx+1, xticks=[], yticks=[])
    pyplot.imshow(numpy.transpose(images[idx], (1, 2, 0)))
    ax.set_title(classes[labels[idx]])

vgg16 = models.vgg16(pretrained=True)

print(vgg16)

print(vgg16.classifier[6].in_features) 
print(vgg16.classifier[6].out_features) 

for param in vgg16.features.parameters():
    param.requires_grad = False

classifier = nn.Sequential(OrderedDict([
    ("Fullly Connected Classifier", nn.Linear(in_features=4096, out_features=CLASS_COUNT, bias=True)),
]))
vgg16.classifier[6] = classifier

vgg16.to(device)

n_epochs = EPOCHS = 2
def train(model: nn.Module, epochs: int=EPOCHS, model_number: int=0,
          epoch_offset: int=1, print_every: int=10) -> tuple:
    """Train, validate, and save the model
    This trains the model and validates it, saving the best 
    (based on validation loss) as =model_<number>_cifar.pth=

    Args:
     model: the network to train
     epochs: number of times to repeat training
     model_number: an identifier for the saved hyperparameters file
     epoch_offset: amount of epochs that have occurred previously
     print_every: how often to print output
    Returns:
     filename, training-loss, validation-loss, improvements: the outcomes for the training
    """
    optimizer = optimize.SGD(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    output_file = "model_{}_vgg.pth".format(model_number)
    training_losses = []
    validation_losses = []
    improvements = []
    valid_loss_min = numpy.Inf # track change in validation loss
    epoch_start = epoch_offset
    last_epoch = epoch_start + epochs + 1
    for epoch in range(epoch_start, last_epoch):

        # keep track of training and validation loss
        train_loss = 0.0
        valid_loss = 0.0

        model.train()
        for data, target in train_loader:
            # move tensors to GPU if CUDA is available            
            data, target = data.to(device), target.to(device)
            # clear the gradients of all optimized variables
            optimizer.zero_grad()
            # forward pass: compute predicted outputs by passing inputs to the model
            output = model(data)
            # calculate the batch loss
            loss = criterion(output, target)
            # backward pass: compute gradient of the loss with respect to model parameters
            loss.backward()
            # perform a single optimization step (parameter update)
            optimizer.step()
            # update training loss
            train_loss += loss.item() * data.size(0)

        model.eval()
        for data, target in valid_loader:
            # move tensors to GPU if CUDA is available
            data, target = data.to(device), target.to(device)
            # forward pass: compute predicted outputs by passing inputs to the model
            output = model(data)
            # calculate the batch loss
            loss = criterion(output, target)
            # update total validation loss 
            valid_loss += loss.item() * data.size(0)

        # calculate average losses
        train_loss = train_loss/len(train_loader.dataset)
        valid_loss = valid_loss/len(valid_loader.dataset)

        # print training/validation statistics 
        if not (epoch % print_every):
            print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(
                epoch, train_loss, valid_loss))
        training_losses.append(train_loss)
        validation_losses.append(valid_loss)
        # save model if validation loss has decreased
        if valid_loss <= valid_loss_min:
            print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(
            valid_loss_min,
            valid_loss))
            torch.save(model.state_dict(), output_file)
            valid_loss_min = valid_loss
            improvements.append(epoch - 1)
    return output_file, training_losses, validation_losses, improvements

def test(best_model):
    criterion = nn.CrossEntropyLoss()
    # track test loss
    test_loss = 0.0
    class_correct = list(0. for i in range(10))
    class_total = list(0. for i in range(10))

    best_model.to(device)
    best_model.eval()
    # iterate over test data
    for data, target in test_loader:
        # move tensors to GPU if CUDA is available
        data, target = data.to(device), target.to(device)
        # forward pass: compute predicted outputs by passing inputs to the model
        output = best_model(data)
        # calculate the batch loss
        loss = criterion(output, target)
        # update test loss 
        test_loss += loss.item() * data.size(0)
        # convert output probabilities to predicted class
        _, pred = torch.max(output, 1)    
        # compare predictions to true label
        correct_tensor = pred.eq(target.data.view_as(pred))
        correct = (
            numpy.squeeze(correct_tensor.numpy())
            if not train_on_gpu
            else numpy.squeeze(correct_tensor.cpu().numpy()))
        # calculate test accuracy for each object class
        for i in range(BATCH_SIZE):
            label = target.data[i]
            class_correct[label] += correct[i].item()
            class_total[label] += 1

    # average test loss
    test_loss = test_loss/len(test_loader.dataset)
    print('Test Loss: {:.6f}\n'.format(test_loss))

    for i in range(10):
        if class_total[i] > 0:
            print('Test Accuracy of %5s: %2d%% (%2d/%2d)' % (
                classes[i], 100 * class_correct[i] / class_total[i],
                numpy.sum(class_correct[i]), numpy.sum(class_total[i])))
        else:
            print('Test Accuracy of %5s: N/A (no training examples)' % (classes[i]))

    print('\nTest Accuracy (Overall): %2d%% (%2d/%2d)' % (
        100. * numpy.sum(class_correct) / numpy.sum(class_total),
        numpy.sum(class_correct), numpy.sum(class_total)))

output_file, training_losses, validation_losses, improvements = train(vgg16, print_every=1)

training_losses = []
validation_losses = []
improvements = []
valid_loss_min = numpy.Inf # track change in validation loss
for epoch in range(1, 3):

    # keep track of training and validation loss
    train_loss = 0.0
    valid_loss = 0.0

    vgg16.train()
    for data, target in train_loader:
        # move tensors to GPU if CUDA is available            
        data, target = data.to(device), target.to(device)
        # clear the gradients of all optimized variables
        optimizer.zero_grad()
        # forward pass: compute predicted outputs by passing inputs to the model
        output = vgg16(data)
        # calculate the batch loss
        loss = criterion(output, target)
        # backward pass: compute gradient of the loss with respect to model parameters
        loss.backward()
        # perform a single optimization step (parameter update)
        optimizer.step()
        # update training loss
        train_loss += loss.item() * data.size(0)

    vgg16.eval()
    for data, target in valid_loader:
        # move tensors to GPU if CUDA is available
        data, target = data.to(device), target.to(device)
        # forward pass: compute predicted outputs by passing inputs to the model
        output = vgg16(data)
        # calculate the batch loss
        loss = criterion(output, target)
        # update total validation loss 
        valid_loss += loss.item() * data.size(0)

    # calculate average losses
    train_loss = train_loss/len(train_loader.dataset)
    valid_loss = valid_loss/len(valid_loader.dataset)

    # print training/validation statistics 
    print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'.format(
        epoch, train_loss, valid_loss))
    training_losses.append(train_loss)
    validation_losses.append(valid_loss)
    # save model if validation loss has decreased
    if valid_loss <= valid_loss_min:
        print('Validation loss decreased ({:.6f} --> {:.6f}).  Saving model ...'.format(
        valid_loss_min,
        valid_loss))
        torch.save(vgg16.state_dict(), output_file)
        valid_loss_min = valid_loss
        improvements.append(epoch - 1)