Generic trainer for Pytorch model

trainers.py

import torch
import numpy as np

class Trainer:

    ##########################
    #                        #
    #     Initialization     #
    #                        #
    ##########################



    def __init__(self,
                 model,
                 dataloader,
                 loss_function,
                 optimizer = None,
                 lr = 0.0005,
                 device='cuda'):
        

        
        ## Assign class attributes and send model to device
        self.model = model.to(device)
        self.dataloader = dataloader
        
        
        ## Initialize the optimizer if none has been passed in
        if optimizer is None:
            self.optimizer = torch.optim.Adam(model.parameters(), lr = lr)
        else:
            self.optimzer = optimizer


        ## Initialize the loss function(s)
        self.loss_function = loss_function
        
        ## Bookkeeping
        self.device = device  
        self.curr_epoch = 0

    ##########################
    #                        #
    #  Single Iter Training  #
    #                        #
    ##########################


    def train_iter(self, input, label, verbose=0):

        ## Zero the gradients
        self.optimizer.zero_grad()

        ## Pass the inputs through the model
        out = self.model(*input)

        ## Calculate the loss(es)
        loss = self.loss_function(out, label)


        ## Pass the loss backward
        loss.backward()

        ## Take an optimizer step
        self.optimizer.step()

        ## Return the total loss
        return(loss)

    ##########################
    #                        #
    #  Mutlti Epoch Training #
    #                        #
    ##########################



    def train(self,
              epochs,
              print_every=1,
              writer=None):

        
        ## Loop over epochs in the range of epochs
        for epoch in range(self.curr_epoch, self.curr_epoch + epochs):
            
            epoch_losses = []

            ## If the report_every epoch is reached, reinitialize metric lists
            if epoch % print_every == 0:
                print("----- Epoch: " + str(epoch) + " -----")
                batch_losses = []
    

            ## Enumerate self.dataloader
            for idx, data_dict in enumerate(self.dataloader):

                ## Grab an example
                x = data_dict["x"]; y = data_dict["y"]

                ## Send it to self.device
                x = x.to(self.device); y = y.to(self.device)
                
                ## Try to train_iter
                batch_loss = self.train_iter(x, y)

                ## Update the metric lists and counters
                batch_losses.append(batch_loss.item())
            
            epoch_losses.append(np.mean(batch_losses))
                
            self.curr_epoch += 1
            
    
            ## If we've hit report_every epoch, print the report
            if epoch % print_every == 0:
                print("avg train loss: " + str(np.mean(epoch_losses)))


            ## Write the outputs to Tensorboard if writer is not None