BaseNetTorch.py

import torch
from torch import nn
import numpy as np
from config import config 
import logging 
from torch.utils.tensorboard import SummaryWriter

from torch_models.torch_utils.training import train_loop, test_loop

class Prediction_history:
    """
    Collect predictions of the given validation set after each epoch 
    predhis is a list of lists (one for each epoch) of tensors (one for each batch)
    """
    def __init__(self, dataloader) -> None:
        self.dataloader = dataloader
        self.predhis = []

    def on_epoch_end(self, model):
        y_pred = []
        for x, y in self.dataloader:
            y_pred.append(model(x.float()))
        self.predhis.append(y_pred)

class BaseNet(nn.Module):
    """
    BaseNet class for ConvNet and EEGnet to inherit common functionality 
    """
    def __init__(self, input_shape, epochs=50, verbose=True, model_number=0, batch_size=64):
        """
        Initialize common variables of models based on BaseNet
        Create the common output layer dependent on the task to run 
        """
        super().__init__()
        self.input_shape = input_shape
        self.epochs = epochs
        self.verbose = verbose
        self.model_number = model_number
        self.batch_size = batch_size
        self.nb_channels = self.input_shape[1]
        self.timesamples = self.input_shape[0]

        # Create output layer depending on task and
        if config['task'] == 'prosaccade_clf':
            self.loss_fn = nn.BCELoss()
            self.output_layer = nn.Sequential(
                nn.Linear(in_features=self.get_nb_features_output_layer(), out_features=1),
                nn.Sigmoid()
            )
        elif config['task'] == 'gaze-reg':
            self.loss_fn = nn.MSELoss()
            self.output_layer = nn.Sequential(
                nn.Linear(in_features=self.get_nb_features_output_layer(), out_features=2) 
            )
        else: #elif config['task'] == 'angle-reg':
            from torch_models.torch_utils.custom_losses import angle_loss
            self.loss_fn = angle_loss
            self.output_layer = nn.Sequential(
                nn.Linear(in_features=self.get_nb_features_output_layer(), out_features=1) 
            )

    # abstract method 
    def forward(self, x):
        """
        Implements a forward pass of the network 
        """
        pass
    
    # abstract method 
    def get_nb_features_output_layer(self):
        """
        Return the number of features that the output layer should take as input
        """
        pass

    # abstract method
    def _split_model(self):
        pass
        
    def fit(self, train_dataloader, test_dataloader, subjectID=None):
        """
        Fit the model on the dataset defined by data x and labels y 
        """
        print("------------------------------------------------------------------------------------")
        print(f"Fitting model number {self.model_number}")
        
        # Create the optimizer
        optimizer = torch.optim.Adam(list(self.parameters()), lr=config['learning_rate'])
        # Create a history to track ensemble performance 
        prediction_ensemble = Prediction_history(dataloader=test_dataloader)
        # Create a summary writer for logging metrics 
        writer = SummaryWriter(log_dir=config['model_dir']+'/summary_writer')
        # Train the model 
        epochs = config['epochs']             
        for t in range(epochs):
            print(f"Epoch {t+1}\n-------------------------------")
            # Run through training and test set 
            train_loss = train_loop(train_dataloader, self.float(), self.loss_fn, optimizer)
            test_loss, test_acc = test_loop(test_dataloader, self.float(), self.loss_fn)
            # Add the predictions on the validation set
            prediction_ensemble.on_epoch_end(model=self)
            # Log metrics to the writer 
            writer.add_scalar('Loss/train', train_loss, t)
            writer.add_scalar('Loss/test', test_loss, t)
            if config['task'] == 'prosaccade-clf':
                writer.add_scalar('Accuracy/test', test_acc, t)
        print(f"Finished model number {self.model_number}")
        if config['save_model']:
            ckpt_dir = config['model_dir'] + '/best_models/' + config['model'] + '_nb_{}_'.format(self.model_number) + 'best_model.pth'
            torch.save(self.state_dict(), ckpt_dir)
        return prediction_ensemble