Trainer.py

import sys
import os
import copy
import logging
from time import time

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
from torchvision import models
from torch.utils.tensorboard import SummaryWriter
from matplotlib import pyplot as plt
from data.data import create_dataset
from seed import seed_everything
from DropLR import DropLR
from options.train_options import TrainOptions


TORCH_SEED = 0

class Trainer:

    def __init__(self, train_opt, dataloaders, dataset_sizes, device, logger=None):
        self.opt = train_opt
        self.dataloaders = dataloaders
        self.dataset_sizes = dataset_sizes
        self.device = device
        self.logger = logger

    def lr_continue_training(self, scheduler, lr_limit=1e-6):
        current_lr = scheduler.get_last_lr()[0]
        if current_lr <= lr_limit:
            return False
        else:
            return True

    def save_model(self, epoch, model_state, optimizer_state, path, loss=None):
        
        save_dir = f'{self.opt.checkpoints_dir}/{self.opt.name}'
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        torch.save({
                'epoch': epoch,
                'model_state_dict': model_state,
                'optimizer_state_dict': optimizer_state,
                'loss': loss,
                }, path)
        
    def log_metrics(self, epoch, phase, logs):
        self.logger.info(f"[{phase}] Epoch {epoch}: loss = {logs['loss']:.4f}, accuracy = {logs['accuracy']:.4f}")
        
    def train_model(self, model, m, criterion, optimizer, scheduler, writer=None, num_epochs=25):
        epoch = -1
        best_model_weights = copy.deepcopy(model.state_dict())
        best_accuracy = 0.0
        minimum_loss = -1

        early_stopping = DropLR(tolerance=self.opt.earlystop_epoch)
        stop_training = False
        drop_learning_rate = False

        if self.opt.continue_train:
            if int(self.opt.epoch) == -1:
                model_path = f'{self.opt.checkpoints_dir}/{self.opt.name}/best_model.pth'
                self.logger.info('Loading the best model and resuming training.')
            else:
                model_path = f'{self.opt.checkpoints_dir}/{self.opt.name}/model_{self.opt.epoch}.pth'
                self.logger.info(f'Loading the model from epoch {self.opt.epoch}.')
                
            checkpoint = torch.load(model_path)
            model.load_state_dict(checkpoint['model_state_dict'])
            epoch = checkpoint['epoch']
            self.logger.info(f'Training resumed after epoch {epoch}')

        for epoch in range(epoch + 1, num_epochs):

            self.logger.info(f"[==== STARTING EPOCH = {epoch} ====]")

            for phase in ['train', 'val']:
                if phase == "train":
                    model.train()
                else:
                    model.eval()
            
                running_loss = 0.0
                running_corrects = 0

                for inputs, labels in self.dataloaders[phase]:
                    
                    inputs = inputs.to(self.device)
                    labels = labels.to(self.device)

                    with torch.set_grad_enabled(phase == "train"):
                        
                        outputs = model(inputs)
                        preds = torch.round(m(outputs))

                        labels = torch.unsqueeze(labels, 1).float()

                        loss = criterion(outputs, labels)

                        if phase == "train":
                            optimizer.zero_grad()
                            loss.backward()
                            optimizer.step()

                    # statistics
                    running_loss += loss.item()
                    running_corrects += torch.sum(preds == labels.data)

                epoch_loss = running_loss / self.dataset_sizes[phase]
                epoch_acc = running_corrects.double() / self.dataset_sizes[phase]

                logs = {'loss': epoch_loss, 'accuracy': epoch_acc}
                self.log_metrics(epoch, phase, logs)

                if writer:
                    writer.add_scalar(f"Loss/{phase}", epoch_loss, epoch)
                    writer.add_scalar(f"Accuracy/{phase}", epoch_acc, epoch)

                if phase == "train" and drop_learning_rate:
                    scheduler.step()
                    self.logger.info(f'Learning rate dropped: {scheduler.get_last_lr()}')

                    if not self.lr_continue_training(scheduler):
                        self.logger.info('Training process stopped. Learning rate dropped below the limit.')
                        stop_training = True
                        break
                
                if phase == "val":
                    drop_learning_rate = early_stopping(epoch_acc)

                if epoch % self.opt.save_epoch_freq == 0:
                    model_path = f'{self.opt.checkpoints_dir}/{self.opt.name}/model_{epoch}.pth'
                    self.save_model(epoch, model.state_dict(), optimizer.state_dict(), model_path, loss)
                
                # deep copy the model
                if phase == 'val' and epoch_acc > best_accuracy:
                    minimum_loss = epoch_loss
                    best_accuracy = epoch_acc
                    best_model_weights = copy.deepcopy(model.state_dict())
                    best_model_optimizer = copy.deepcopy(optimizer.state_dict())

                    best_model_path = f'{self.opt.checkpoints_dir}/{self.opt.name}/best_model.pth'
                    self.save_model(epoch, best_model_weights, best_model_optimizer, best_model_path)

            if stop_training:
                break
            
        model.load_state_dict(best_model_weights)
        if writer:
            writer.flush()
            writer.close()

        return model

if __name__=="__main__":

    seed_everything(TORCH_SEED)

    train_opt = TrainOptions().parse()
    val_opt = TrainOptions().parse()
    val_opt.isTrain = False

    train_logger = logging.getLogger("Train_Logger")
    train_logger.setLevel(logging.INFO)
    train_fh = logging.FileHandler(filename=f'{train_opt.logs_dir}/{train_opt.log_file}.log', encoding='utf-8')        
    train_fh.setLevel(logging.INFO)
    train_logger.addHandler(train_fh)

    opts = {'train': train_opt, 'val': val_opt}

    writer = SummaryWriter(f'runs/{train_opt.name}/train')

    data_directory = "dataset"
    datasets = ["train", "val"]
    dataloaders = {}
    dataset_sizes = {}

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    classes = train_opt.classes
    train_logger.info(f'Classes: {classes}')

    for dataset in datasets:
        dst_start = time()

        dataloaders[dataset], dataset_sizes[dataset] = create_dataset(f"{data_directory}/{dataset}", opts[dataset], opts[dataset].classes)
        train_logger.info(f'[==== DATASET CREATION DURATION FOR {dataset} = {time() - dst_start} ====]')

    # Import the pretrained ResNet50 Model trained on the ImageNet Dataset
    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
    num_features = model.fc.in_features

    model.fc = nn.Linear(num_features, 1)
    
    model = model.to(device)

    m = nn.Sigmoid()
    criterion = nn.BCEWithLogitsLoss()
    
    if train_opt.optim == "adam":
        optimizer = optim.Adam(model.parameters(), lr=train_opt.lr)
    elif train_opt.optim == "sgd":
        optimizer = optim.SGD(model.parameters(), lr=train_opt.lr)
    else:
        train_logger.info('Optimizer not valid!')
        sys.exit()

    step_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
    
    trainer = Trainer(train_opt, dataloaders, dataset_sizes, device, train_logger)

    model = trainer.train_model(model, m, criterion, optimizer, step_lr_scheduler, writer)