train_SSUG.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-
# @Time    :  2022-12-2 21:45
# @Author  : He Li
"""
Apply feature disentanglement in multi sites fundus segmentation
"""
import argparse
import logging
import os
import tqdm
import random
import timeit
import torch
import numpy as np
import torch.optim as optim
import matplotlib.pyplot as plt
from torchvision import transforms
from torch import nn
import torch, torchvision
import torch.nn.functional as F

import supervision
from dataloader.fundus_dataloader import FundusSegmentation
from dataloader.fourier_augmentation import LogRadialEMASTATS, SimpleLogRadialExtrapCHW
from torchvision.utils import make_grid
from torch.utils.data import ConcatDataset, DataLoader
from models.weight_init import initialize_weights
from dataloader import fundus_transforms as tr
from utils.average_meter import AverageMeter
import utils
from utils.losses import *
from tensorboardX import SummaryWriter
from pytorch_metric_learning import losses, distances, reducers, regularizers, miners

# -------------------- reproduction ------------------------ #
# torch.cuda.current_device()  # no problem in server
# torch.cuda._initialized = True  # no problem in server
torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.cuda.manual_seed_all(42)
np.random.seed(42)  # Numpy module.
random.seed(42)  # Python random module.
torch.manual_seed(42)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.enabled = True
torch.set_default_tensor_type('torch.FloatTensor')

@torch.no_grad()
def _pca_rgb_single(feat_chw: torch.Tensor) -> torch.Tensor:
    """Single feature PCA→RGB, feat:[C,H,W](cpu,f32) -> [3,H,W]∈[0,1]"""
    C,H,W = feat_chw.shape
    X = feat_chw.view(C, -1).t()          # [HW,C]
    X = X - X.mean(dim=0, keepdim=True)
    U,S,Vh = torch.linalg.svd(X, full_matrices=False)
    P = (U[:, :3] * S[:3]).view(H, W, 3)  # [H,W,3]
    P = (P - P.amin((0,1), True)) / (P.amax((0,1), True) - P.amin((0,1), True) + 1e-6)
    
    return P.permute(2,0,1).contiguous()  # [3,H,W]

@torch.no_grad()
def log_embedding_compact(writer, emb_img, emb_FDA, images, epoch, step, args,
                          tag_prefix="Train/EmbCompact", viz_num=4, up_to_input=True):
    """
    Only two visualizations are recorded:
    1) PCA pseudocolor: img/FDA
    2) Global cosine similarity curve: img-FDA
    """
    B, C, H, W = emb_img.shape
    n = min(B, viz_num)

    # ---Prepare CPU float32 copy (save video memory/avoid AMP)---
    e_img = emb_img[:n].detach().to('cpu', dtype=torch.float32)  # [n,C,H,W]
    e_fda = emb_FDA[:n].detach().to('cpu', dtype=torch.float32)

    # ---------- (1) PCA pseudocolor ----------
    pca_imgs, pca_fdas = [], []
    upsz = (args.img_size, args.img_size) if up_to_input else (H, W)
    for i in range(n):
        p3_img = _pca_rgb_single(e_img[i])
        p3_fda = _pca_rgb_single(e_fda[i])
        if up_to_input:
            p3_img = F.interpolate(p3_img.unsqueeze(0), size=upsz, mode='nearest').squeeze(0)
            p3_fda = F.interpolate(p3_fda.unsqueeze(0), size=upsz, mode='nearest').squeeze(0)
        pca_imgs.append(p3_img); pca_fdas.append(p3_fda)

    writer.add_image(f'{tag_prefix}/PCA_img',
                     torchvision.utils.make_grid(torch.stack(pca_imgs, 0), nrow=n), step)
    writer.add_image(f'{tag_prefix}/PCA_FDA',
                     torchvision.utils.make_grid(torch.stack(pca_fdas, 0), nrow=n), step)

    # ---------- (2) Global cosine similarity curve ----------
    def _avg_pair_cos(a_bchw, b_bchw):
        a = a_bchw.mean(dim=(2,3))  # [B,C]
        b = b_bchw.mean(dim=(2,3))
        a = F.normalize(a, dim=1); b = F.normalize(b, dim=1)
        return (a * b).sum(dim=1).mean().item()

    gcos_if = _avg_pair_cos(e_img, e_fda)

    writer.add_scalars(f'{tag_prefix}/global_cosine',
                       {'img-FDA': gcos_if}, step)


def _init_gate_tracker():
    return {
        'count': 0,
        'block_gate_sum': None,       # [L]
        'channel_supp_sum': None,     # [L, C]
    }


@torch.no_grad()
def _update_gate_tracker(tracker, block_gate_mean, block_channel_gate_mean):
    block_gate_mean = block_gate_mean.detach().float().cpu()
    channel_supp = (1.0 - block_channel_gate_mean.detach().float()).cpu()
    if tracker['block_gate_sum'] is None:
        tracker['block_gate_sum'] = torch.zeros_like(block_gate_mean)
        tracker['channel_supp_sum'] = torch.zeros_like(channel_supp)
    tracker['block_gate_sum'] += block_gate_mean
    tracker['channel_supp_sum'] += channel_supp
    tracker['count'] += 1


def _finalize_gate_tracker(tracker):
    if tracker['count'] == 0:
        return None
    avg_block_gate = tracker['block_gate_sum'] / tracker['count']     # [L]
    avg_block_supp = 1.0 - avg_block_gate                             # [L]
    avg_channel_supp = tracker['channel_supp_sum'] / tracker['count'] # [L, C]
    global_channel_supp = avg_channel_supp.mean(dim=0)                # [C]
    return {
        'avg_block_gate': avg_block_gate,
        'avg_block_supp': avg_block_supp,
        'avg_channel_supp': avg_channel_supp,
        'global_channel_supp': global_channel_supp,
    }


def _topk_channels(channel_supp, topk):
    k = min(topk, channel_supp.shape[-1])
    values, indices = torch.topk(channel_supp, k=k, dim=-1)
    return indices, values


def _format_topk(indices, values):
    return ', '.join([f'{int(i)}:{float(v):.4f}' for i, v in zip(indices.tolist(), values.tolist())])


def _save_gate_stats_npz(save_dir, epoch, domain_id, gate_stats, topk, save_full_matrix=False):
    os.makedirs(save_dir, exist_ok=True)
    blk_top_idx, blk_top_val = _topk_channels(gate_stats['avg_channel_supp'], topk=topk)  # [L, K]
    g_top_idx, g_top_val = _topk_channels(gate_stats['global_channel_supp'].unsqueeze(0), topk=topk)  # [1, K]

    payload = {
        'avg_block_gate': gate_stats['avg_block_gate'].numpy(),
        'avg_block_supp': gate_stats['avg_block_supp'].numpy(),
        'global_channel_supp': gate_stats['global_channel_supp'].numpy(),
        'block_topk_indices': blk_top_idx.numpy(),
        'block_topk_supp': blk_top_val.numpy(),
        'global_topk_indices': g_top_idx.squeeze(0).numpy(),
        'global_topk_supp': g_top_val.squeeze(0).numpy(),
    }
    if save_full_matrix:
        payload['avg_channel_supp'] = gate_stats['avg_channel_supp'].numpy()

    save_path = os.path.join(save_dir, f'epoch_{epoch+1:04d}_domain{domain_id}_gate_stats.npz')
    np.savez_compressed(save_path, **payload)
    return save_path


def parse_args():
    desc = "Pytorch implementation of SSUG (HeLi)"
    parser = argparse.ArgumentParser(description=desc)
    # dir config
    parser.add_argument('--exp_dir', type=str, default='./exp/Proposed/SSUG')
    parser.add_argument('--data_dir', type=str, default='./dataset/Fundus')
    parser.add_argument('--workspace', type=str, default='./exp/Proposed/SSUG/checkpoint')
    # GPU config
    parser.add_argument('--gpu', type=str, default='0')
    parser.add_argument('--gpu_grop', type=int, default=[0, 1])
    # training config
    parser.add_argument('--resume', default=None, help='checkpoint path')
    parser.add_argument('--datasetTrain', nargs='+', type=int, default=[1],
                        help='train folder id contain images ROIs to train range from [1,2,3,4]')
    parser.add_argument('--datasetTest', nargs='+', type=int, default=[1],
                        help='test folder id contain images ROIs to test one of [1,2,3,4]')
    parser.add_argument('--in_channel', type=int, default=3)
    parser.add_argument('--batch_size', type=int, default=8)
    parser.add_argument('--num_classes', type=int, default=2)
    parser.add_argument('--con_w', type=float, default=0.1)
    parser.add_argument('--vit_name', type=str, default='vit_b', help='select one vit model')
    parser.add_argument('--ckpt', type=str, default='./checkpoint/sam_vit_b_01ec64.pth', help='Pretrained checkpoint')
    parser.add_argument('--img_size', type=int, default=512, help='input patch size of network input')
    parser.add_argument('--start_epoch', type=int, default=0)
    parser.add_argument('--epoches', type=int, default=400)
    parser.add_argument('--learning_rate', type=float, default=0.0001)
    parser.add_argument('-wi', '--weight_init', type=str, default="xavier",
                        help='Weight initialization method, or path to weights file '
                             '(for fine-tuning or continuing training)')
    parser.add_argument('--print_interval', type=int, default=1)
    parser.add_argument('--val_interval', type=int, default=2)
    parser.add_argument('--save_interval', type=int, default=25)
    parser.add_argument('--gate_topk', type=int, default=16)
    parser.add_argument('--gate_save_interval', type=int, default=20)
    parser.add_argument('--gate_save_matrix', action='store_true',
                        help='save full [num_blocks, num_channels] suppression matrix to npz')

    # utils.check_folder(parser.parse_args().exp_dir)
    return parser.parse_args()


def validate_slice(model, dataloader, img_folder, anatomy_folder, args, writer, epoch, multimask_output):
    training = model.training
    model.eval()
    
    domain_dice = []
    for i in range(len(args.datasetTest)):
        val_dice_domain = AverageMeter()
        domain_dice.append({'disesc': val_dice_domain})
    with torch.no_grad():
        for sample in tqdm.tqdm(dataloader, total=len(dataloader), ncols=80, leave=False):
            for dc, domain in enumerate(sample):
                img, gt = domain['image'].cuda(), domain['label'].cuda()
                
                outputs = model(img, multimask_output, args.img_size)
                seg_pred = outputs['masks']
                seg_pred = torch.sigmoid(seg_pred)
                domain_dice[dc]['disesc'].update(val_dice_class(seg_pred.permute(0, 2, 3, 1) > 0.75, gt.permute(0, 2, 3, 1),
                                               num_class=args.num_classes))

                if (epoch + 1) % (args.print_interval + 19) == 0:
                    grid_image = make_grid(img, nrow=args.batch_size, normalize=True)
                    writer.add_image('Valid/Domain{}/imgs'.format(args.datasetTest[dc]), grid_image, epoch)
                    label_shape = gt.size()
                    grid_image = make_grid(
                        gt.reshape(label_shape[0] * label_shape[1], label_shape[2], label_shape[3]).unsqueeze(dim=1),
                        nrow=args.batch_size, normalize=True)
                    writer.add_image('Valid/Domain{}/mask'.format(args.datasetTest[dc]), grid_image, epoch)
                    pred_shape = seg_pred.size()
                    grid_image = make_grid(
                        seg_pred.reshape(pred_shape[0] * pred_shape[1], pred_shape[2], pred_shape[3]).unsqueeze(dim=1),
                        nrow=args.batch_size, normalize=True)
                    writer.add_image('Valid/Domain{}/prediction'.format(args.datasetTest[dc]), grid_image, epoch)
    if training:
        model.train()

    return domain_dice


def train(model, train_loader, val_loader, writer, args, multimask_output):
    
    # SME
    ema = LogRadialEMASTATS(C=args.in_channel, R=256, momentum=0.99, device=torch.device('cuda'))
    Freq_Aug = SimpleLogRadialExtrapCHW(
        ema,
        beta_range=(0.22, 0.45),
        smooth_band=0.25,
        gamma_range=(0.8, 1.5)
    )

    # define the criterion
    dice_criterion = DiceLoss()
    bce_criterion = torch.nn.BCEWithLogitsLoss()
    
    # Auto Mix Precision
    scaler = torch.cuda.amp.GradScaler()
    
    # define the optimizer
    optimizer = optim.Adam(model.parameters(), betas=(0.9, 0.99), lr=args.learning_rate)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.95, patience=4, verbose=True,
                                                     min_lr=1e-6)
    os.makedirs(args.workspace, exist_ok=True)
    gate_stat_dir = os.path.join(args.workspace, 'gate_stats')
    os.makedirs(gate_stat_dir, exist_ok=True)

    best_domain_dice = []
    best_domain_epoch = []
    for i in range(len(args.datasetTest)):
        best_val_dice, best_epoch = torch.tensor([0.0, 0.0]), 1
        best_domain_dice.append(best_val_dice)
        best_domain_epoch.append(best_epoch)

    for epoch in range(args.start_epoch, args.epoches):
        domain_loss = []
        for i in range(len(args.datasetTrain)):
            seg_loss_epoch = AverageMeter()
            fda_loss_epoch = AverageMeter()
            content_loss_epoch = AverageMeter()
            gate_score_epoch = AverageMeter()
            gate_tracker = _init_gate_tracker()
            domain_loss.append({'FDA': fda_loss_epoch, 'con': content_loss_epoch, 'seg': seg_loss_epoch,
                                'gate': gate_score_epoch, 'gate_tracker': gate_tracker})

        # train in each epoch
        start_time = timeit.default_timer()
        for batch_idx, sample in tqdm.tqdm(
                enumerate(train_loader), total=len(train_loader),
                desc='Train epoch=%d' % epoch, ncols=80, leave=False):
            iteration = batch_idx + epoch * len(train_loader)

            model.train()
            for dc, domain in enumerate(sample):
                image = domain['image'].cuda()
                # GLA_img = domain['GLA_image'].cuda()
                # FDA_img = domain['FDA_image'].cuda()
                label = domain['label'].cuda().float()
                
                use_aug = ema.initialized
                if use_aug:
                    FDA_img = torch.stack([Freq_Aug(image[i]) for i in range(image.size(0))], 0)
                else:
                    FDA_img = image
                ema.update(image)
                
                optimizer.zero_grad(set_to_none=True)
                
                with torch.cuda.amp.autocast(enabled=False):
                    # UACG paired forward (init/freq) with symmetric channel calibration
                    paired_out = model.forward_train_uacg(image, FDA_img, multimask_output, args.img_size)
                    init_out = paired_out['init_outputs']
                    freq_out = paired_out['freq_outputs']

                    # FDA branch after UACG
                    emb_FDA = freq_out['img_embeddings']
                    seg_FDA = freq_out['masks']
                    bce_FDA = bce_criterion(seg_FDA, label)

                    # init branch after UACG
                    emb_img = init_out['img_embeddings']
                    seg_img = init_out['masks']
                    prob_img = torch.sigmoid(seg_img)
                    dice = dice_criterion(prob_img.permute(0,2,3,1), label.permute(0,2,3,1), num_class=args.num_classes)
                    bce_img = bce_criterion(seg_img, label)
                    seg_loss = 0.5*(dice + bce_img)
                    
                    # UACG consistency on calibrated image embeddings
                    content_loss = (
                        torch.nn.functional.l1_loss(emb_img, emb_FDA, reduction='mean') +
                        torch.nn.functional.l1_loss(emb_FDA, emb_img, reduction='mean'))
                    
                    # UACG consistency on calibrated intermediate features
                    medfeat_content_loss = paired_out['uacg_con_loss']
                    content_loss += medfeat_content_loss

                    gate_mean = paired_out['uacg_gate_mean']
                    gate_block_mean = paired_out['uacg_gate_block_mean']
                    gate_channel_mean = paired_out['uacg_gate_channel_mean']
                    
                    loss = bce_FDA + seg_loss + args.con_w * content_loss
                
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

                scaler.step(optimizer)
                scaler.update()
                
                domain_loss[dc]['FDA'].update(bce_FDA.detach().cpu())
                domain_loss[dc]['seg'].update(seg_loss.detach().cpu())
                domain_loss[dc]['con'].update(content_loss.detach().cpu())
                domain_loss[dc]['gate'].update(gate_mean.detach().cpu())
                _update_gate_tracker(domain_loss[dc]['gate_tracker'], gate_block_mean, gate_channel_mean)
                
                # tensorboard for visualing train result
                if (epoch + 1) % (args.print_interval + 19) == 0:
                    grid_image = make_grid(image, nrow=args.batch_size, normalize=True)
                    writer.add_image('Train/Init_imgs', grid_image, epoch)
                    grid_image = make_grid(FDA_img, nrow=args.batch_size, normalize=True)
                    writer.add_image('Train/FDA_imgs', grid_image, epoch)
                    label_shape = label.size()
                    grid_image = make_grid(
                        label.reshape(label_shape[0] * label_shape[1], label_shape[2], label_shape[3]).unsqueeze(dim=1),
                        nrow=args.batch_size, normalize=True)
                    writer.add_image('Train/mask', grid_image, epoch)
                    pred_shape = seg_img.size()
                    grid_image = make_grid(
                        seg_img.reshape(pred_shape[0] * pred_shape[1], pred_shape[2], pred_shape[3]).unsqueeze(dim=1),
                        nrow=args.batch_size, normalize=True)
                    writer.add_image('Train/prediction', grid_image, epoch)
                    
                    log_embedding_compact(
                        writer=writer,
                        emb_img=emb_img, emb_FDA=emb_FDA,
                        images=image,
                        epoch=epoch, step=epoch, args=args,
                        tag_prefix="Train/EmbCompact", viz_num=min(4, args.batch_size),
                        up_to_input=True
                    )

        # print training result
        logging.info(
            '\n Epoch[%4d/%4d]-Lr: %.6f --> Train...' % (epoch + 1, args.epoches, optimizer.param_groups[0]['lr']))
        for i in range(len(args.datasetTrain)):
            logging.info(
                '\t Domain-%d : Seg Loss = %.4f, FDA Seg Loss = %.4f, Con Loss = %.4f, Gate Mean = %.4f' %
                (args.datasetTrain[i], domain_loss[i]['seg'].avg, domain_loss[i]['FDA'].avg,
                 domain_loss[i]['con'].avg, domain_loss[i]['gate'].avg))

        epoch_gate_stats = []
        for i, train_dc in enumerate(args.datasetTrain):
            gate_stats = _finalize_gate_tracker(domain_loss[i]['gate_tracker'])
            epoch_gate_stats.append(gate_stats)
            if gate_stats is None:
                continue

            g_top_idx, g_top_val = _topk_channels(gate_stats['global_channel_supp'].unsqueeze(0), args.gate_topk)
            g_top_idx, g_top_val = g_top_idx.squeeze(0), g_top_val.squeeze(0)
            logging.info('\t Domain-%d : Global Top-%d Suppressed Channels -> %s' %
                         (train_dc, min(args.gate_topk, gate_stats['global_channel_supp'].numel()),
                          _format_topk(g_top_idx, g_top_val)))
            blk_top_idx, blk_top_val = _topk_channels(gate_stats['avg_channel_supp'], args.gate_topk)
            for blk_id in range(blk_top_idx.shape[0]):
                logging.info('\t   Block-%02d Top-%d Suppressed -> %s' %
                             (blk_id, blk_top_idx.shape[1], _format_topk(blk_top_idx[blk_id], blk_top_val[blk_id])))

            if (epoch + 1) % args.gate_save_interval == 0:
                save_path = _save_gate_stats_npz(
                    save_dir=gate_stat_dir,
                    epoch=epoch,
                    domain_id=train_dc,
                    gate_stats=gate_stats,
                    topk=args.gate_topk,
                    save_full_matrix=args.gate_save_matrix,
                )
                logging.info('\t Domain-%d : gate stats saved to %s' % (train_dc, save_path))

        # tensorboard
        if (epoch + 1) % args.print_interval == 0:
            writer.add_scalar('Learning_Rate', optimizer.param_groups[0]['lr'], epoch)
            for i, train_dc in enumerate(args.datasetTrain):
                writer.add_scalars('Train/Domain{}/Losses'.format(train_dc),
                                   {'Seg': domain_loss[i]['seg'].avg, 'FDA': domain_loss[i]['FDA'].avg,
                                    'Con': domain_loss[i]['con'].avg, 'GateMean': domain_loss[i]['gate'].avg}, epoch)
                gate_stats = epoch_gate_stats[i]
                if gate_stats is None:
                    continue

                writer.add_scalars('Train/Domain{}/UACG_BlockGate'.format(train_dc),
                                   {f'blk_{b:02d}': gate_stats['avg_block_gate'][b].item()
                                    for b in range(gate_stats['avg_block_gate'].numel())}, epoch)
                writer.add_scalars('Train/Domain{}/UACG_BlockSuppression'.format(train_dc),
                                   {f'blk_{b:02d}': gate_stats['avg_block_supp'][b].item()
                                    for b in range(gate_stats['avg_block_supp'].numel())}, epoch)
                writer.add_histogram('Train/Domain{}/UACG_GlobalChannelSuppression'.format(train_dc),
                                     gate_stats['global_channel_supp'], epoch)

                g_top_idx, g_top_val = _topk_channels(gate_stats['global_channel_supp'].unsqueeze(0), args.gate_topk)
                topk_text = _format_topk(g_top_idx.squeeze(0), g_top_val.squeeze(0))
                writer.add_text('Train/Domain{}/UACG_GlobalTopSuppressedChannels'.format(train_dc), topk_text, epoch)

        # validate and visualization
        if not os.path.exists(args.workspace):
            os.mkdir(args.workspace)
        result_dir = os.path.join(args.workspace, 'val_results')
        if not os.path.exists(result_dir):
            os.mkdir(result_dir)
        model_dir = os.path.join(args.workspace, 'models')
        if not os.path.exists(model_dir):
            os.mkdir(model_dir)
        model_both_domain_dir = []
        for test_dc in args.datasetTest:
            model_domain_dir = os.path.join(model_dir, 'Domain%d' % test_dc)
            if not os.path.exists(model_domain_dir):
                os.mkdir(model_domain_dir)
            model_both_domain_dir.append(model_domain_dir)
        if (epoch + 1) % args.val_interval == 0:
            val_img_path = os.path.join(result_dir, 'Ep_%04d_imgs.png' % (epoch + 1))
            val_anatomy_path = os.path.join(result_dir, 'Ep_%04d_anatomys.png' % (epoch + 1))
            domain_dice = validate_slice(model, val_loader, val_img_path, val_anatomy_path, args, writer, epoch, multimask_output)
            
            logging.info('\n Epoch[%4d/%4d] --> Valid...' % (epoch + 1, args.epoches))
            for i in range(len(args.datasetTest)):
                logging.info(
                    '\t Domain-%d: [Dice Coef: mean=%.4f, cup=%.4f, disc=%.4f]' %
                    (args.datasetTest[i], torch.mean(domain_dice[i]['disesc'].avg), domain_dice[i]['disesc'].avg[0], domain_dice[i]['disesc'].avg[1]))
            for i, test_dc in enumerate(args.datasetTest):
                writer.add_scalars('Val/Domain{}/Dice'.format(test_dc),
                                   {'cup': domain_dice[i]['disesc'].avg[0], 'disc': domain_dice[i]['disesc'].avg[1],
                                    'mean': torch.mean(domain_dice[i]['disesc'].avg)}, epoch)
            # save best model
            for i in range(len(args.datasetTest)):
                if torch.mean(domain_dice[i]['disesc'].avg) >= torch.mean(best_domain_dice[i]):
                    best_model_path = os.path.join(model_both_domain_dir[i], 'best_model.pth')
                    torch.save(model.state_dict(), best_model_path)
                    logging.info(
                        '\n Domain-%d: [Epoch[%4d/%4d] --> Dice improved from %.4f (cup=%.4f, disc=%.4f in epoch %4d) '
                        'to %.4f (cup=%.4f, disc=%.4f)]' %
                        (args.datasetTest[i], epoch + 1, args.epoches, torch.mean(best_domain_dice[i]), best_domain_dice[i][0], best_domain_dice[i][1],
                         best_domain_epoch[i], torch.mean(domain_dice[i]['disesc'].avg), domain_dice[i]['disesc'].avg[0], domain_dice[i]['disesc'].avg[1]))
                    best_domain_dice[i], best_domain_epoch[i] = domain_dice[i]['disesc'].avg, epoch + 1
                else:
                    logging.info('\n Domain-%d: [Epoch[%4d/%4d] --> Dice did not improved with %.4f (cup=%.4f, '
                                 'disc=%.4f in epoch %d)]' %
                                 (args.datasetTest[i], epoch + 1, args.epoches, torch.mean(best_domain_dice[i]), best_domain_dice[i][0], best_domain_dice[i][1],
                                  best_domain_epoch[i]))
            # check for plateau
            dice_sum = 0
            for i in range(len(args.datasetTest)):
                dice_sum += torch.mean(domain_dice[i]['disesc'].avg)
            avg = dice_sum/len(args.datasetTest)
            scheduler.step(avg)

        # save model
        if (epoch + 1) >= (args.epoches - 100) and (epoch + 1) % args.save_interval == 0:
            for i in range(len(args.datasetTest)):
                model_path = os.path.join(model_both_domain_dir[i], 'Ep_%04d_dice_%.4f.pth' % ((epoch + 1), torch.mean(domain_dice[i]['disesc'].avg)))
                torch.save({'model': model.state_dict(), 'optim': optimizer.state_dict()}, model_path)
                logging.info('\t Domain-%d: [Save Model] to %s' % (args.datasetTest[i], model_path))


def main():
    args = parse_args()

    # GPU Setting
    # single GPU
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    # many GPU
    # device = torch.device("cuda:%d" %(args.gpu_grop[0]) if torch.cuda.is_available() else "cpu")

    # define logger
    logging.basicConfig(filename=os.path.join(args.exp_dir, 'train.log'), level=logging.DEBUG,
                        format='%(asctime)s %(message)s')
    logging.getLogger().addHandler(logging.StreamHandler())

    # print all parameters
    for name, v in vars(args).items():
        logging.info(name + ': ' + str(v))

    # dataset
    composed_transforms_tr = transforms.Compose([
        tr.RandomScaleCrop(512),
        tr.RandomRotate(),
        tr.RandomFlip(),
        tr.elastic_transform(),
        tr.add_salt_pepper_noise(),
        tr.adjust_light(),
        tr.eraser(),
        tr.One_hot(),
        tr.ToTensor()
    ])

    composed_transforms_ts = transforms.Compose([
        tr.RandomCrop(512),
        tr.Normalize_tf(),
        tr.ToTensor()
    ])

    # dataloader config
    train_set = FundusSegmentation(base_dir=args.data_dir, phase='train', splitid=args.datasetTrain,
                                   transform=composed_transforms_tr)
    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=6, drop_last=True,
                              pin_memory=True)
    valid_set = FundusSegmentation(base_dir=args.data_dir, phase='test', splitid=args.datasetTest,
                                   transform=composed_transforms_ts)
    valid_loader = DataLoader(valid_set, batch_size=1, shuffle=False, num_workers=6, pin_memory=True)

    #  single GPU
    from models.SSUG.segment_anything import sam_model_registry
    sam, img_embedding_size = sam_model_registry[args.vit_name](image_size=args.img_size,
                                                                num_classes=args.num_classes,
                                                                checkpoint=args.ckpt, pixel_mean=[0, 0, 0],
                                                                pixel_std=[1, 1, 1])
    model = sam.cuda()
    low_res = img_embedding_size * 4
    
    if args.num_classes > 1:
        multimask_output = True
    else:
        multimask_output = False

    print('parameter numer:', sum([p.numel() for p in model.parameters()]))

    # GPU Parallel
    # if torch.cuda.device_count() > 1:
    #     model = DataParallel(model, device_ids=args.gpu_grop).cuda()
    # model.to(device)

    if args.resume:
        checkpoint = torch.load(args.resume)
        pretrained_dict = checkpoint['model']
        model_dict = model.state_dict()
        # 1. filter out unnecessary keys
        pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
        # 2. overwrite entries in the existing state dict
        model_dict.update(pretrained_dict)
        # 3. load the new state dict
        model.load_state_dict(model_dict)
        print('Resume finished!')
    else:
        initialize_weights(model, args.weight_init)

    # summary writer config
    writer = SummaryWriter(log_dir=args.exp_dir, comment=args.exp_dir.split('/')[-1])

    # train
    train(model, train_loader, valid_loader, writer, args, multimask_output)


if __name__ == '__main__':
    main()