model_helpers.py

import torch
import torch.nn as nn
import os
import pickle
from sklearn.metrics import confusion_matrix

train_params = {
	'batch_size': 256,
    'shuffle': True
}

val_params = {
	'batch_size': 256,
    'shuffle': False
}

test_params = {
	'batch_size': 256,
    'shuffle': False
}

num_classes = 8
image_size = 96

def do_setup(USE_GPU, type):
	global dtype
	global device
	
	dtype = type

	if USE_GPU and torch.cuda.is_available(): 
		print("Using gpu")
		device = torch.device('cuda')
	else:
		print("Using cpu")
		device = torch.device('cpu')



def check_accuracy(loader, model):
    """
    Finds the accuracy of a model
    
    Inputs:
    - loader: A dataloader containing the validation / testing dataset
    - model: A PyTorch Module giving the model to evaluate.
    
    Returns: 
    - The number of correct predictions
    - The number of samples
    - The accuracy of the model
    """

    num_correct = 0
    num_samples = 0
    model = model.to(device=device)  # move the model parameters to CPU/GPU
    model.eval()  # set model to evaluation mode
    with torch.no_grad(): # no need to store computation graph or local gradients
        for x, y in loader:
            x = x.to(device=device, dtype=dtype)  # move to device, e.g. GPU
            y = y.to(device=device, dtype=torch.long)
            scores = model(x)
            _, preds = scores.max(1)
            num_correct += (preds == y).sum()
            num_samples += preds.size(0)
        acc = float(num_correct) / num_samples
        return num_correct, num_samples, 100 * acc
    

def train(model, optimizer, loader_train, loader_val, epochs=1, scheduler=None, print_every=1):
    """
    Train a model using the PyTorch Module API.
    
    Inputs:
    - model: A PyTorch Module giving the model to train.
    - optimizer: An Optimizer object we will use to train the model
    - loader_train: A dataloader containing the train dataset
    - loader_val: A dataloader containing the validation dataset
    - epochs: (Optional) An integer giving the number of epochs to train for
    - scheduler: (Optional) A learning rate scheduler
    - print_every: (Optional) An integer specifying how often to print the loss. 
    
    Prints model losses and accuracies during training.

    Returns: 
    - The training losses after each epoch
    - The validation accuracy after each epoch
    """
    training_losses = []
    training_accuracies = []
    validation_losses = []
    validation_accuracies = []
    model = model.to(device=device)  # move the model parameters to CPU/GPU
    for e in range(epochs):
        total_training_loss = 0
        num_iterations_training = 0

        total_validation_loss = 0
        num_iterations_validation = 0

        for t, (x, y) in enumerate(loader_train):
            model.train()  # put model to training mode
            x = x.to(device=device, dtype=dtype)  # move to device, e.g. GPU
            y = y.to(device=device, dtype=torch.long)

            scores = model(x)
            loss = torch.nn.functional.cross_entropy(scores, y)

            # Since we are updating the model as we go along, this
            # only estimates the loss of the whole epoch
            total_training_loss += loss.item()
            num_iterations_training += 1

            # Zero out all of the gradients for the variables which the optimizer
            # will update.
            optimizer.zero_grad()

            # This is the backwards pass: compute the gradient of the loss with
            # respect to each  parameter of the model.
            loss.backward()

            # Actually update the parameters of the model using the gradients
            # computed by the backwards pass.
            optimizer.step()

            # if t % print_every == 0:
            #     print('Epoch {}, iteration {}, loss = {}'.format(e, t, loss.item()))
        
        # Calculate the validation loss
        for t, (x, y) in enumerate(loader_val):
            model.eval()  # set model to evaluation mode
            x = x.to(device=device, dtype=dtype)  # move to device, e.g. GPU
            y = y.to(device=device, dtype=torch.long)

            scores = model(x)
            loss = torch.nn.functional.cross_entropy(scores, y)
            total_validation_loss += loss.item()
            num_iterations_validation += 1
        
        training_losses.append(total_training_loss / num_iterations_training)
        validation_losses.append(total_validation_loss / num_iterations_validation)
        
        # Calculate accuracies
        training_num_correct, training_num_samples, training_accuracy = check_accuracy(loader_train, model)
        training_accuracies.append(training_accuracy)
        validation_num_correct, validation_num_samples, validation_accuracy = check_accuracy(loader_val, model)
        validation_accuracies.append(validation_accuracy)

        if type(scheduler) is torch.optim.lr_scheduler.ReduceLROnPlateau:
            scheduler.step(validation_accuracy)
        elif scheduler is not None:
            scheduler.step()

        print((
              "Epoch %d: "
              "training loss = %f, training accuracy = %d / %d correct (%.2f), "
              "validation loss = %f, validation accuracy = %d / %d correct (%.2f)"
              ) %
              (e,
               training_losses[-1], training_num_correct, training_num_samples, training_accuracies[-1],
               validation_losses[-1], validation_num_correct, validation_num_samples, validation_accuracies[-1]))  

        # print("Epoch %d: training loss = %f, validation accuracy = %d / %d correct (%.2f)" %
        #       (e, total_training_loss / num_iterations_training, validation_num_correct, validation_num_samples, validation_accuracy))

     
    return (training_losses, training_accuracies, validation_losses, validation_accuracies)



def check_accuracy_for_ensemble(model_info):
    """
    Finds the accuracy of a model
    
    Inputs:
    - loader: A dataloader containing the validation / testing dataset
    - model: A PyTorch Module giving the model to evaluate.
    
    Returns: 
    - The number of correct predictions
    - The number of samples
    - The accuracy of the model
    """

    num_correct = 0
    num_samples = 0
        
    with torch.no_grad(): # no need to store computation graph or local gradients
        final_results = []
        for (loader, model) in model_info:
            model = model.to(device=device)
            model.eval()  # set model to evaluation mode
            results = []
            for (x, y) in loader:
                x = x.to(device=device, dtype=dtype)  # move to device, e.g. GPU
                y = y.to(device=device, dtype=torch.long)
                results.append(model(x))
            final_results.append(results)
        
        all_preds = []

        for i in range(len(final_results[0])):
            scores = torch.sum(torch.stack([results[i] for results in final_results]), 0)
            _, preds = scores.max(1)
            all_preds.append(preds)
        
        for t, (x, y) in enumerate(model_info[0][0]):
            y = y.to(device=device, dtype=torch.long)
            num_correct += (all_preds[t] == y).sum()
            num_samples += all_preds[t].size(0)
        
        acc = float(num_correct) / num_samples
        return num_correct, num_samples, 100 * acc, all_preds


def create_confusion_matrix(model, loader):
	true_labels = []
	pred_labels = []
	model = model.to(device=device)
	model.eval()  # set model to evaluation mode
	with torch.no_grad(): # no need to store computation graph or local gradients
		for x, y in loader:
			x = x.to(device=device, dtype=dtype)  # move to device, e.g. GPU
			y = y.to(device=device, dtype=torch.long)
			scores = model(x)
			_, preds = scores.max(1)

			pred_labels.extend(preds.cpu().numpy())
			true_labels.extend(y.cpu().numpy())
    
	return confusion_matrix(true_labels, pred_labels, normalize='true')


def load_model(model, model_file):
	model_file = "results/" + model_file
	if os.path.exists(model_file):
		model.load_state_dict(torch.load(model_file))

def save_model(model, model_file):
	model_file = "results/" + model_file
	if not os.path.exists(model_file):
		torch.save(model.state_dict(), model_file)

def load_results(filename):
	filename = "results/" + filename
	with open(filename, 'rb') as f:
		return pickle.load(f)

def save_results(results, filename):
	filename = "results/" + filename
	with open(filename, 'wb') as f:
		pickle.dump(results, f)


def flatten(x):
    N = x.shape[0] 
    return x.view(N, -1)  

class Flatten(nn.Module):
    def forward(self, x):
        return flatten(x)