model.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Author: Jan Ruhland


import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torchvision.models as models
from tqdm import tqdm
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau
from torch.utils.data import DataLoader
import os
from sklearn.metrics import accuracy_score
from preprocessing import Preprocessor
from torch.cuda.amp import autocast, GradScaler

class CustomResNetModel(nn.Module):
    def __init__(self, num_additional_features=446, freeze=True, tag='left'):
        super(CustomResNetModel, self).__init__()
        
        self.tag = tag
        self.num_additional_features = num_additional_features
        if self.tag == 'left':
            self.num_additional_features = 222
        elif self.tag == 'right':
            self.num_additional_features = 225

        # Load pre-trained ResNet50 model
        resnet = models.resnet18(pretrained=True)

        if freeze:
            # Freeze the pre-trained layers
            for param in resnet.parameters():
                param.requires_grad = False
        
        # Remove the final classification layer
        self.resnet = nn.Sequential(*list(resnet.children())[:-1])
        
        # Add custom layers for regression (age prediction)
        self.regression_age = nn.Sequential(
            nn.Linear(resnet.fc.in_features + self.num_additional_features, 512),
            nn.LeakyReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 1)
        )

        # Add custom layers for binary classification (normal or not)
        self.classification_binary = nn.Sequential(
            nn.Linear(resnet.fc.in_features + self.num_additional_features, 512),
            nn.LeakyReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 1),
            # nn.Sigmoid()  # Sigmoid activation for binary classification
        )

    def forward(self, image, additional_features):
        # Forward pass through the pre-trained ResNet
        x = self.resnet(image)
        
        # Global average pooling
        x = x.view(x.size(0), -1)
        
        # Concatenate additional features
        x_combined = torch.cat([x, additional_features], dim=1)
        
        # Custom layers for regression (age prediction)
        age_output = self.regression_age(x_combined)

        # Custom layers for binary classification (normal or not)
        binary_output = self.classification_binary(x_combined)

        return age_output, binary_output
    
    
def finalPred(image, additional_features=None, checkpoint=None, tag='left'):
    """
    The model takes an image as input and returns the prediction.

    Parameters
    ----------
    image : path to image
        DESCRIPTION.
    Returns
    -------
    age.

    """
    # TODO: Take the input image and return a predicted age.
    
    # Load the model
    model = CustomResNetModel(tag=tag)

    # Load checkpoint if provided
    if checkpoint is not None:
        try:
            dict_stat = torch.load(checkpoint, map_location='cpu')
            if tag == 'left':
                model.load_state_dict(dict_stat['left_model_state_dict'])
            elif tag == 'right':
                model.load_state_dict(dict_stat['right_model_state_dict'])
        except FileNotFoundError:
            print("Warning: Checkpoint file not found.")


    # additional_features might be None so we need a default value of an empty tensor with the correct shape 454
    if additional_features is None:
        if tag == 'left':
            additional_features = torch.zeros((1, 222), dtype=torch.float32)
        elif tag == 'right':
            additional_features = torch.zeros((1, 225), dtype=torch.float32)

    # Load the checkpoint
    # if checkpoint is not None:
    #     dict_stat = torch.load(checkpoint, map_location=device)
    #     left_model.load_state_dict(dict_stat['left_model_state_dict'])
    #     right_model.load_state_dict(dict_stat['right_model_state_dict'])

    # Set the model to evaluation mode
    model.eval()

    # Preprocess the image
    image = Preprocessor(image)

    # Make a prediction
    with torch.no_grad():
        age, binary = model(image, additional_features)

    binary = torch.sigmoid(binary)
    binary = torch.round(binary).int()

    # Return the predicted age
    # return (age.item(), binary.item())
    return age.item()


# Function to train the model for one epoch
def train_one_epoch(model, dataloader, optimizer, device='cpu'):
    model = model.to(device)
    model.train()
    scaler = GradScaler()
    total_loss = 0.0
    age_criterion = nn.L1Loss()
    # binary_criterion = nn.BCELoss()
    positive_weight = torch.tensor([2]).to(device)
    binary_criterion = nn.BCEWithLogitsLoss(pos_weight=positive_weight)

    for batch_idx, (image, additional_features, age, binary_output) in enumerate(tqdm(dataloader, desc='Training')):
        # Move data to device
        image, additional_features, age, binary_output = (
            image.to(device),
            additional_features.to(device),
            age.to(device),
            binary_output.to(device),
        )

        # Zero the gradients
        optimizer.zero_grad()

        # Forward pass with autocast for mixed precision
        with autocast():
            # Forward pass
            age_out, n_out = model(image, additional_features)

            # Apply sigmoid activation for the binary output
            # n_out = torch.sigmoid(n_out)

            # Compute the loss
            loss_age = age_criterion(age_out, age.float().view(-1, 1))
            loss_binary = binary_criterion(n_out, binary_output.float().view(-1, 1))
            loss = loss_age + loss_binary

        # Backward pass using GradScaler
        scaler.scale(loss).backward()

        # Update weights using GradScaler
        scaler.step(optimizer)

        # Update GradScaler for the next iteration
        scaler.update()

        # Accumulate total loss
        total_loss += loss.cpu().item()

    # Return average loss
    return total_loss / len(dataloader)


# Function to evaluate the model on the validation set
def evaluate_model(model, dataloader, device='cpu'):
    model = model.to(device)
    model.eval()
    total_loss = 0.0
    mae = 0.0
    correct_binary = 0
    age_criterion = nn.L1Loss()
    # binary_criterion = nn.BCELoss()
    binary_criterion = nn.BCEWithLogitsLoss()

    with torch.no_grad():
        for image, additional_features, age, binary_target in tqdm(dataloader, desc='Evaluating'):
            # Move data to device
            image, additional_features, age, binary_target = (
                image.to(device),
                additional_features.to(device),
                age.to(device),
                binary_target.to(device),
            )

            # Forward pass
            age_out, binary_output = model(image, additional_features)
            loss_age = age_criterion(age_out, age.float().view(-1, 1))
            # apply sigmoid activation for the binary output
            binary_prediction = torch.sigmoid(binary_output)
            loss_binary = binary_criterion(binary_output, binary_target.float().view(-1, 1))
            loss = loss_age + loss_binary


            # Accumulate total loss
            total_loss += loss.cpu().detach()

            # Calculate MAE
            mae += loss_age.cpu().detach()

            # Calculate accuracy for binary classification
            binary_prediction = torch.round(binary_prediction).int()
            correct_binary += accuracy_score(binary_prediction.cpu(), binary_target.cpu(), normalize=False)

    # Return average loss and accuracy
    return total_loss / len(dataloader), mae / len(dataloader), correct_binary / len(dataloader.dataset)


def train_model(model, train_dataset, test_dataset, lr=1e-3, weight_decay=1e-5, batch_size=128,
                device='cpu', num_epochs=5, tag=None, checkpoint=None, early_stopping_patience=5, print_plot=True):
    
    # Early stopping variables
    best_val_loss = float('inf')
    early_stopping_counter = 0
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=weight_decay)
    scheduler = ReduceLROnPlateau(optimizer, mode='min', patience=early_stopping_patience-1, factor=0.5, verbose=True)
    train_losses = []
    val_losses = []

    # Instantiate data loaders
    if device != 'cpu':
        pin_memory=True
    else:
        pin_memory=False
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=pin_memory)
    val_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=pin_memory)

    # Load checkpoint if provided
    if checkpoint is not None:
        try:
            dict_stat = torch.load(checkpoint, map_location=device)
            if tag == 'left':
                model.load_state_dict(dict_stat['left_model_state_dict'])
                optimizer.load_state_dict(dict_stat['left_optimizer_state_dict'])
            elif tag == 'right':
                model.load_state_dict(dict_stat['right_model_state_dict'])
                optimizer.load_state_dict(dict_stat['right_optimizer_state_dict'])
        except FileNotFoundError:
            print("Warning: Checkpoint file not found. Training from scratch.")

    # Training loop
    for epoch in range(num_epochs):
        # Training
        train_loss = train_one_epoch(model, train_loader, optimizer, device)

        # Validation
        val_loss, mae, accuracy = evaluate_model(model, val_dataloader, device)

        # Print or log training and validation metrics
        print(f'Epoch {epoch + 1}:' +
            f' Train Loss: {train_loss:.4f}' +
            f' Val Loss: {val_loss:.4f}' +
            f' MAE: {mae:.4f}' +
            f' Accuracy: {accuracy:.4f}')

        # Append losses and accuracy for plotting
        train_losses.append(train_loss)
        val_losses.append(val_loss)

        # Learning rate scheduler step
        scheduler.step(val_loss)

        # Early stopping
        if train_loss < best_val_loss:
            best_val_loss = train_loss
            early_stopping_counter = 0
        else:
            early_stopping_counter += 1
            if early_stopping_counter >= early_stopping_patience:
                print("Early stopping!")
                break

    # Plot the training and validation losses
    if print_plot:
        plt.figure(figsize=(10, 5))
        plt.plot(train_losses, label='Train Loss')
        plt.plot(val_losses, label='Validation Loss')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.title('Training and Validation Losses')
        plt.legend()
        plt.show()

    # save the model and optimizer states for later use
    # trainable_params_dict = {name: param for name, param in model.cpu().named_parameters() if param.requires_grad}
    if checkpoint is not None:
        torch.save({
            tag + '_model_state_dict': model.cpu().state_dict(),
            tag + '_optimizer_state_dict': optimizer.state_dict(),
        }, checkpoint)
    else:
        # check if the checkpoint folder exists
        if not os.path.exists('checkpoint'):
            os.makedirs('checkpoint')

        checkpoint = 'checkpoint/left_model.pt' if tag == 'left' else 'checkpoint/right_model.pt'
        torch.save({
            tag + '_model_state_dict': model.cpu().state_dict(),
            tag + '_optimizer_state_dict': optimizer.state_dict(),
        }, checkpoint)
        
    return model