train.py

"""
@author: Originally by Francesco La Rosa
         Adapted by Vatsal Raina, Nataliia Molchanova
"""

import argparse
import os
import torch
from torch import nn
from monai.data import decollate_batch
from monai.transforms import Compose, AsDiscrete
from monai.inferers import sliding_window_inference
from monai.losses import DiceLoss
from monai.metrics import DiceMetric
from monai.utils import set_determinism
import numpy as np
import random
from metrics import dice_metric
from data_load import get_train_dataloader, get_val_dataloader
import wandb
from os.path import join as pjoin
from metrics import *
from model import *
import time

parser = argparse.ArgumentParser(description='Get all command line arguments.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
# trainining
parser.add_argument('--frozen_learning_rate', type=float, default=-1, help='Specify the initial learning rate')
parser.add_argument('--learning_rate', type=float, default=1e-5, help='Specify the initial learning rate')
parser.add_argument('--n_epochs', type=int, default=300, help='Specify the number of epochs to train for')
parser.add_argument('--path_model', type=str, default=None, help='Path to pretrained model')
# initialisation                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       
parser.add_argument('--seed', type=int, default=1, help='Specify the global random seed')
# data
parser.add_argument('--data_dir', type=str, default='/data/', help='Specify the path to the data files directory')

parser.add_argument('--I', nargs='+', default=['FLAIR'], choices=['FLAIR', 'phase_T2starw', 'mag_T2starw',\
                                                                  'MPRAGE_reg-T2starw_T1w', 'T1map', 'UNIT1'])

parser.add_argument('--segmentation_mode', type=str, default='semantic', choices=['semantic', 'instances'])
parser.add_argument('--include_ctrl', action='store_true', default=False, help='Include control lesions')
parser.add_argument('--apply_mask', type=str, default=None, 
        help='The name of the folder containing the masks you want to the images to be applied to.')
parser.add_argument('--cp_factor', type=int, default=0, 
			help="number of times each object has been copied for the copy paste data augmentation strategy")

parser.add_argument('--save_path', type=str, default='/models/',
                    help='Specify the path to the save directory')

parser.add_argument('--num_workers', type=int, default=12, help='Number of workers')
parser.add_argument('--batch_size', default=4, type=int)
parser.add_argument('--cache_rate', default=1.0, type=float)
# logging
parser.add_argument('--val_interval', type=int, default=5, help='Validation every n-th epochs')
parser.add_argument('--threshold', type=float, default=0.4, help='Probability threshold')

parser.add_argument('--wandb_project', type=str, default='mslis', help='wandb project name')
parser.add_argument('--name', default="idiot without a name", help='Wandb run name')

VAL_AMP = True
roi_size = (96, 96, 96)

def check_paths(args):
    from os.path import exists as pexists
    assert pexists(args.data_dir), f"Directory not found {args.data_dir}"
    #assert pexists(args.bm_path), f"Directory not found {args.bm_path}"
    assert pexists(args.save_path), f"Directory not found {args.save_path}"
    if args.path_model:
        assert pexists(args.path_model), f"Warning: File not found {args.path_model}"


def inference(input, model):
    def _compute(input):
        return sliding_window_inference(
            inputs=input,
            roi_size=roi_size,
            sw_batch_size=2,
            predictor=model,
            overlap=0.5,
        )

    if VAL_AMP:
        with torch.cuda.amp.autocast():
            return _compute(input)
    else:
        return _compute(input)


def get_default_device():
    """ Set device """
    if torch.cuda.is_available():
        print("Got CUDA!")
        return torch.device('cuda')
    else:
        raise Exception("CUDA WAS NOT DETECTED !!!")
        return torch.device('cpu')


post_trans = Compose(
    [AsDiscrete(argmax=True, to_onehot=2)]
)


def main(args):
    check_paths(args)
    seed_val = args.seed
    random.seed(seed_val)
    np.random.seed(seed_val)
    torch.manual_seed(seed_val)
    torch.cuda.manual_seed_all(seed_val)
    set_determinism(seed=seed_val)

    device = get_default_device()
    torch.multiprocessing.set_sharing_strategy('file_system')

    save_dir = f'{args.save_path}/{args.name}'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    wandb.login()
    wandb.init(project=args.wandb_project, mode="online")
    wandb.run.name = f'{args.name}'
    '''' Initialize dataloaders '''
    
    train_loader = get_train_dataloader(data_dir=args.data_dir, 
                         num_workers= args.num_workers, 
                         I=args.I, 
                         mode=args.segmentation_mode, 
                         include_ctrl=bool(args.include_ctrl),
                         cache_rate=args.cache_rate,
                         apply_mask=args.apply_mask,
			 cp_factor=args.cp_factor)
    val_loader = get_val_dataloader(data_dir=args.data_dir, 
                                    num_workers= args.num_workers,
                                    I=args.I, 
                                    mode=args.segmentation_mode, 
                                    include_ctrl=bool(args.include_ctrl),
                                    cache_rate=args.cache_rate,
                                    apply_mask=args.apply_mask)

    ''' Initialise the model '''
    if args.path_model is not None:
        print(f"Retrieving pretrained model from {args.path_model}")
        model = get_pretrained_model(args.path_model, len(args.I)).to(device)
    else:
        print(f"Initializing new model with {len(args.I)} input channels")
        model = UNet3D(in_channels=len(args.I), num_classes=2).to(device)

    #print(model)

    loss_function = DiceLoss(to_onehot_y=True,
                             softmax=True, sigmoid=False,
                             include_background=False)
    # Separate out the first layer parameters
    first_layer_params = model.a_block1.conv1.parameters()
    
    # Separate out the rest of the parameters
    rest_of_model_params = [p for p in model.parameters() if p not in first_layer_params]
    flr = args.learning_rate if args.frozen_learning_rate < 0 else args.frozen_learning_rate
    
    # Initialize the optimizer with two parameter groups having different learning rates
    optimizer = torch.optim.Adam([{'params': first_layer_params, 'lr': args.learning_rate},
                                  {'params': rest_of_model_params, 'lr': flr}],
                                 weight_decay=0.0005) #momentum=0.9,
    
    #optimizer = torch.optim.Adam(model.parameters(), args.learning_rate, weight_decay=0.0005)
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.n_epochs)#, eta_min=1e-6)

    act = nn.Softmax(dim=1)

    epoch_num = args.n_epochs
    val_interval = args.val_interval
    threshold = args.threshold
    gamma_focal = 2.0
    dice_weight = 0.5
    focal_weight = 1.0

    dice_metric = DiceMetric(include_background=False, reduction="mean")
    best_metric_nDSC, best_metric_epoch_nDSC = -1, -1
    best_metric_DSC, best_metric_epoch_DSC = -1, -1

    epoch_loss_values, metric_values_nDSC, metric_values_DSC = [], [], []

    scaler = torch.cuda.amp.GradScaler()
    step_print = 1
    
    ''' Training loop '''
    for epoch in range(epoch_num):
        start_epoch_time = time.time()
        print("-" * 10)
        print(f"epoch {epoch + 1}/{epoch_num}")
        model.train()
        epoch_loss = 0
        epoch_loss_ce = 0
        epoch_loss_dice = 0
        step = 0
        for batch_data in train_loader:
            n_samples = batch_data["image"].size(0)
            for m in range(0, batch_data["image"].size(0), args.batch_size):
                step += args.batch_size
                inputs, labels = (
                    batch_data["image"][m:(m + 2)].to(device),
                    batch_data["label"][m:(m + 2)].type(torch.LongTensor).to(device))

                with torch.cuda.amp.autocast():
                    outputs = model(inputs)

                    # Dice loss
                    loss1 = loss_function(outputs, labels)
                    # Focal loss
                    ce_loss = nn.CrossEntropyLoss(reduction='none')
                    ce = ce_loss(outputs, torch.squeeze(labels, dim=1))
                    pt = torch.exp(-ce)
                    loss2 = (1 - pt) ** gamma_focal * ce
                    loss2 = torch.mean(loss2)
                    loss = dice_weight * loss1 + focal_weight * loss2

                epoch_loss += loss.item()
                epoch_loss_ce += loss2.item()
                epoch_loss_dice += loss1.item()

                scaler.scale(loss).backward()

                scaler.step(optimizer)
                scaler.update()

                optimizer.zero_grad()
          
                if step % 100 == 0:
                    elapsed_time = time.time() - start_epoch_time
                    step_print = int(step / args.batch_size)
                    print(
                        f"{step_print}/{(len(train_loader) * n_samples) // (train_loader.batch_size * args.batch_size)}, train_loss: {loss.item():.4f}" + \
                        f"(elapsed time: {int(elapsed_time // 60)}min {int(elapsed_time % 60)}s)") 

        epoch_loss /= step_print
        epoch_loss_dice /= step_print
        epoch_loss_ce /= step_print
        epoch_loss_values.append(epoch_loss)
        print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

        current_lr = optimizer.param_groups[0]['lr']
        lr_scheduler.step()

        wandb.log(
            {'Total Loss/train': epoch_loss, 'Dice Loss/train': epoch_loss_dice, 'Focal Loss/train': epoch_loss_ce,
             'Learning rate': current_lr, },  # 'Dice Metric/train': metric},
            step=epoch)
        ''' Validation '''
        if (epoch + 1) % val_interval == 0 and epoch > 5:
            model.eval()
            with torch.no_grad():
                nDSC_list = []
                for val_data in val_loader:
                    val_inputs, val_labels, val_bms = (
                        val_data["image"].to(device),
                        val_data["label"].to(device),
                        val_data["brain_mask"].squeeze().cpu().numpy()
                    )
                    #val_inputs, val_labels = (
                    #    val_data["image"].to(device),
                    #    val_data["label"].to(device),
                    #)

                    val_outputs = inference(val_inputs, model)

                    for_dice_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]

                    dice_metric(y_pred=for_dice_outputs, y=val_labels)

                    val_outputs = act(val_outputs)[:, 1]
                    val_outputs = torch.where(val_outputs >= threshold, torch.tensor(1.0).to(device),
                                              torch.tensor(0.0).to(device))
                    val_outputs = val_outputs.squeeze().cpu().numpy()
                    #curr_preds = thresholded_output.squeeze().cpu().numpy()[val_bms == 1]
                    #gts = val_labels.squeeze().cpu().numpy()[val_bms == 1]
                    #nDSC = dice_norm_metric(gts, curr_preds)
                    nDSC = dice_norm_metric(val_labels.squeeze().cpu().numpy()[val_bms == 1], val_outputs[val_bms == 1])
                    #nDSC = dice_norm_metric(val_labels.squeeze().cpu().numpy(), val_outputs)
                    nDSC_list.append(nDSC)

                torch.cuda.empty_cache()
                del val_inputs, val_labels, val_outputs, for_dice_outputs, val_bms # , thresholded_output, curr_preds, gts , val_bms
                metric_nDSC = np.mean(nDSC_list)
                metric_DSC = dice_metric.aggregate().item()
                wandb.log({'nDSC Metric/val': metric_nDSC, 'DSC Metric/val': metric_DSC}, step=epoch)
                metric_values_nDSC.append(metric_nDSC)
                metric_values_DSC.append(metric_DSC)

                if metric_nDSC > best_metric_nDSC:
                    best_metric_nDSC = metric_nDSC
                    best_metric_epoch_nDSC = epoch + 1
                    save_path = os.path.join(save_dir, f"best_nDSC_{args.name}_seed{args.seed}.pth")
                    torch.save(model.state_dict(), save_path)
                    print("saved new best metric model for nDSC")

                if metric_DSC > best_metric_DSC:
                    best_metric_DSC = metric_DSC
                    best_metric_epoch_DSC = epoch + 1
                    save_path = os.path.join(save_dir, f"best_DSC_{args.name}_seed{args.seed}.pth")
                    torch.save(model.state_dict(), save_path)
                    print("saved new best metric model for DSC")

                print(f"current epoch: {epoch + 1} current mean normalized dice: {metric_nDSC:.4f}"
                      f"\nbest mean normalized dice: {best_metric_nDSC:.4f} at epoch: {best_metric_epoch_nDSC}"
                      f"\nbest mean dice: {best_metric_DSC:.4f} at epoch: {best_metric_epoch_DSC}"
                      )

                dice_metric.reset()


# %%
if __name__ == "__main__":
    args = parser.parse_args()
    main(args)