train_val_swin.py

import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
import numpy as np
from tqdm import tqdm
import os
import argparse
import time

import warnings
warnings.filterwarnings("ignore")

from torch.utils.tensorboard import SummaryWriter

from monai.inferers import sliding_window_inference

from monai.networks.nets import SwinUNETR
from optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR
from types import SimpleNamespace

from gg_tools import dice_score, TEMPLATE, get_key, NUM_CLASS, ORGAN_NAME, get_train_val_data_loader, DiceLoss, Multi_BCELoss


torch.multiprocessing.set_sharing_strategy('file_system')

def aggregate_distributed_losses(args,loss_dice,loss_bce,len_iter):

    total_dice_loss = loss_dice * len_iter
    total_bce_loss = loss_bce * len_iter

    if args.dist:

        tensors = {
            'dice_loss': torch.tensor(total_dice_loss,devcice = args.device),
            'bce_loss': torch.tensor(total_bce_loss,device = args.device),
            'samples': torch.tensor(len_iter,args.device)
        }

        for tensor in tensors.values():
            dist.all_reduce(tensor,op=dist.ReduceOp.SUM)
        
        avg_dice_loss = tensors['dice_loss']/tensors['samples']
        avg_bce_loss = tensors['bce_loss']/tensors['samples']

    else:

        avg_dice_loss = total_dice_loss/len_iter
        avg_bce_loss = total_bce_loss/len_iter
    
    return avg_dice_loss.item(),avg_bce_loss.item()

#training process

def train(args, train_loader, model, optimizer, loss_seg_DICE, loss_seg_CE):

    model.train()

    loss_bce_sum = torch.tensor(0.0).to(args.device)
    loss_dice_sum = torch.tensor(0.0).to(args.device)
    total_steps = torch.tensor(0).to(args.device)

    epoch_iterator = tqdm(
        train_loader,
        desc=f"Epoch {args.epoch}: Training",
        disable=args.local_rank!=0,
        dynamic_ncols=True
    )

    for step, batch in enumerate(epoch_iterator):

        x = batch["image"].to(args.device)
        y = batch["post_label"].to(args.device).float()
        name = batch['name']

        logit_map = model(x)
        term_seg_Dice = loss_seg_DICE(logit_map, y, name, TEMPLATE)
        term_seg_BCE = loss_seg_CE(logit_map, y, name, TEMPLATE)
        loss = term_seg_BCE + term_seg_Dice

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        loss_bce_sum += term_seg_BCE.item()
        loss_dice_sum += term_seg_Dice.item()
        total_steps += 1

        if args.local_rank==0:
            epoch_iterator.set_description(
                f"Epoch {args.epoch}: Training ({step+1}/{len(train_loader)}) "
                f"(dice_loss={term_seg_Dice.item():.5f}, bce_loss={term_seg_BCE.item():.5f})"
            )

        torch.cuda.empty_cache()

    # Synchronize the sum of losses and total steps across all processes
    dist.all_reduce(loss_bce_sum, op=dist.ReduceOp.SUM)
    dist.all_reduce(loss_dice_sum, op=dist.ReduceOp.SUM)
    dist.all_reduce(total_steps, op=dist.ReduceOp.SUM)

    # Calculate average losses
    avg_bce_loss = loss_bce_sum.item() / total_steps.item()
    avg_dice_loss = loss_dice_sum.item() / total_steps.item()

    if args.local_rank==0:
        print(f'Epoch {args.epoch}: avg_dice_loss={avg_dice_loss:.5f}, '
              f'avg_bce_loss={avg_bce_loss:.5f} (across all GPUs)')

    return avg_dice_loss, avg_bce_loss, total_steps.item()

#validation process 
def distributed_validation(model, ValLoader,args):

    model.eval()

    dice_list = {key: torch.zeros(2,NUM_CLASS).to(args.device) for key in TEMPLATE.keys()}

    for batch in tqdm(ValLoader,disable = args.local_rank!=0):

        image, label, name = batch['image'].to(args.device), batch['post_label'], batch['name']

        with torch.no_grad():

            pred = sliding_window_inference(image, (args.roi_x, args.roi_y, args.roi_z), 1, model)
            pred_sigmoid = F.sigmoid(pred)

        B = pred_sigmoid.shape[0]
        for b in range(B):
            template_key = get_key(name[b])
            organ_list = TEMPLATE[template_key]
            for organ in organ_list:
                dice_organ = dice_score(pred_sigmoid[b,organ-1,:,:,:], label[b,organ-1,:,:,:].cuda())[0]
                dice_list[template_key][0][organ-1] += dice_organ.item()
                dice_list[template_key][1][organ-1] += 1

    #accumulate resutls across all GPUs
    for key in TEMPLATE.keys():
        dist.all_reduce(dice_list[key],op = dist.ReduceOp.SUM)

    dice_list = {key: value.cpu().numpy() for key,value in dice_list.items()}

    #calculate average dice scores

    avg_organ_dice = np.zeros((2,NUM_CLASS))

    if args.local_rank == 0:

        with open('out/'+args.log_name+f'/val_{args.epoch}.txt', 'w') as f:
            for key in TEMPLATE.keys():
                organ_list = TEMPLATE[key]
                content = 'Task%s| '%(key)
                for organ in organ_list:
                    dice = dice_list[key][0][organ-1] / dice_list[key][1][organ-1]
                    content += '%s: %.4f, '%(ORGAN_NAME[organ-1], dice)
                    avg_organ_dice[0][organ-1] += dice_list[key][0][organ-1]
                    avg_organ_dice[1][organ-1] += dice_list[key][1][organ-1]
                f.write(content)
                f.write('\n')
            content = 'Average | '
            for i in range(NUM_CLASS):
                content += '%s: %.4f, '%(ORGAN_NAME[i], avg_organ_dice[0][i] / avg_organ_dice[1][i])
            f.write(content)
            f.write('\n')

    return avg_organ_dice


def validation(model, ValLoader, args):

    model.eval()
   
    dice_list = {}

    for key in TEMPLATE.keys():
        dice_list[key] = np.zeros((2, NUM_CLASS)) # 1st row for dice, 2nd row for count

    for index, batch in enumerate(tqdm(ValLoader)):
        image, label, name = batch["image"].cuda(), batch["post_label"], batch["name"]
        
        with torch.no_grad():
            pred = sliding_window_inference(image, (args.roi_x, args.roi_y, args.roi_z), 1, model)
            pred_sigmoid = F.sigmoid(pred)

        B = pred_sigmoid.shape[0]
        for b in range(B):
            template_key = get_key(name[b])
            organ_list = TEMPLATE[template_key]
            for organ in organ_list:
                dice_organ = dice_score(pred_sigmoid[b,organ-1,:,:,:], label[b,organ-1,:,:,:].cuda())[0]
                dice_list[template_key][0][organ-1] += dice_organ.item()
                dice_list[template_key][1][organ-1] += 1
        #if(index == 10): break

    ave_organ_dice = np.zeros((2, NUM_CLASS))

    with open('out/'+args.log_name+f'/val_{args.epoch}.txt', 'w') as f:
        for key in TEMPLATE.keys():
            organ_list = TEMPLATE[key]
            content = 'Task%s| '%(key)
            for organ in organ_list:
                dice = dice_list[key][0][organ-1] / dice_list[key][1][organ-1]
                content += '%s: %.4f, '%(ORGAN_NAME[organ-1], dice)
                ave_organ_dice[0][organ-1] += dice_list[key][0][organ-1]
                ave_organ_dice[1][organ-1] += dice_list[key][1][organ-1]
            f.write(content)
            f.write('\n')
        content = 'Average | '
        for i in range(NUM_CLASS):
            content += '%s: %.4f, '%(ORGAN_NAME[i], ave_organ_dice[0][i] / ave_organ_dice[1][i])
        f.write(content)
        f.write('\n')

    return ave_organ_dice

def process(args):
    
    if args.dist:
        dist.init_process_group(backend="nccl", init_method="env://")
        rank = args.local_rank
        print(args.local_rank)

    args.device = torch.device(f"cuda:{rank}")
    torch.cuda.set_device(args.device)
    print(args.device)

    # Model initialization
    model = SwinUNETR(img_size=(args.roi_x, args.roi_y, args.roi_z),
                            in_channels=1,
                            out_channels=NUM_CLASS,
                            feature_size=48)
    
    if args.pretrain:
        model.load_params(torch.load(args.pretrain)["state_dict"])
    
    model.to(args.device)
    if args.dist:
        model = DDP(model, device_ids=[args.device])
    # criterion and optimizer    
    loss_seg_DICE = DiceLoss(num_classes=NUM_CLASS).to(args.device)
    loss_seg_CE = Multi_BCELoss(num_classes=NUM_CLASS).to(args.device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    scheduler = LinearWarmupCosineAnnealingLR(optimizer, warmup_epochs=args.warmup_epoch, max_epochs=args.max_epoch)

    if args.resume:
        checkpoint = torch.load(args.resume)
        if args.dist:
            model.load_state_dict(checkpoint['net'])
        else:
            store_dict = model.state_dict()
            model_dict = checkpoint['net']
            for key in model_dict.keys():
                store_dict['.'.join(key.split('.')[1:])] = model_dict[key]
            model.load_state_dict(store_dict)
        optimizer.load_state_dict(checkpoint['optimizer'])
        args.epoch = checkpoint['epoch']
        scheduler.load_state_dict(checkpoint['scheduler'])
        
        print('success resume from ', args.resume)

    torch.backends.cudnn.benchmark = True

    train_loader, val_loader,train_sampler, val_sampler = get_train_val_data_loader(args)

    if rank==0:
        writer = SummaryWriter(log_dir='out/' + args.log_name)
        print('Writing Tensorboard logs to ', 'out/' + args.log_name)
        print('training started')
    
    while args.epoch < args.max_epoch:

        if args.dist:
            dist.barrier()
            train_sampler.set_epoch(args.epoch)
            val_sampler.set_epoch(args.epoch)
        
        scheduler.step()

        avg_loss_dice, avg_loss_bce, len_iter = train(args, train_loader, model, optimizer, loss_seg_DICE, loss_seg_CE)
                
        avg_organ_dice_val = distributed_validation(model,val_loader,args) #getting average organ dice loss in validation set.

        if args.local_rank == 0:

            writer.add_scalar('train_dice_loss', avg_loss_dice, args.epoch)
            writer.add_scalar('train_bce_loss', avg_loss_bce, args.epoch)
            writer.add_scalar('lr', np.array(scheduler.get_lr()), args.epoch)        

            for i in range(avg_organ_dice_val.shape[1]):

                writer.add_scalar(f'Dice Score Class {ORGAN_NAME[i]}',avg_organ_dice_val[0][i] / avg_organ_dice_val[1][i],args.epoch)

        if (args.epoch % args.store_num == 0 and args.epoch != 0) and args.local_rank==0:
            checkpoint = {
                "net": model.state_dict(),
                'optimizer':optimizer.state_dict(),
                'scheduler': scheduler.state_dict(),
                "epoch": args.epoch
            }
            if not os.path.isdir('out/' + args.log_name):
                os.mkdir('out/' + args.log_name)
            torch.save(checkpoint, 'out/' + args.log_name + '/epoch_' + str(args.epoch) + '.pth')
            print('save model success')

        args.epoch += 1

def main():

    args = SimpleNamespace(
        space_x = 1.5,
        space_y = 1.5,
        space_z = 1.5,
        roi_x = 96,
        roi_y = 96,
        roi_z = 96,
        num_samples = 2,
        data_root_path = '/blue/kgong/s.kapoor/language_guided_segmentation/CLIP-Driven-Universal-Model_MSD_only/',
        data_txt_path = './dataset/dataset_list/',
        batch_size = 4,
        num_workers = 8,
        a_min = -175,
        a_max = 250,
        b_min = 0.0,
        b_max = 1.0,
        dataset_list = ['PAOT_10_inner'],
        NUM_CLASS = 9,
        backbone = 'swinunetr',
        trans_encoding = 'word_embedding',
        #pretrain = './out/swinunetr_dist_msd/epoch_20.pth',
        pretrain = None,
        lr = 1e-4,
        weight_decay = 1e-5,
        dist = True,
        max_epoch = 500,
        store_num = 10,
        warmup_epoch = 10,
        epoch = 0,
        local_rank = int(os.environ['LOCAL_RANK']),
        device = None,
        resume = None
    )

    #args to parse are as follows:
    parser = argparse.ArgumentParser(description = 'Some arguments to take')
    parser.add_argument('--log_name', default='swinunet', help='The path resume from checkpoint')

    parsed_args = parser.parse_args()

    args_dict = vars(parsed_args)
    for key,value in args_dict.items():
        if value is not None:
            setattr(args,key,value)

    process(args=args)

if __name__ == "__main__":

    main()

# torchrun --nproc_per_node=4 --master_port=1234 train_val_dist.py --dist True --backbone swinunetr --log_name swinunetr_dist_model --batch_size 3