main_train_AMIR.py

import os 
os.environ['CUDA_VISIBLE_DEVICES']='0'  
from model.Model_AMIR import AMIR
from loss.losses import CharbonnierLoss
from evaluation.evaluation_metric import compute_measure
from data.common import transformData, dataIO
from data.MedicalDataUniform import Train_Data, Test_Data, DataSampler
import numpy as np
import pickle
import torch
from torch import nn
from torch.utils.data import DataLoader 
from torch.optim.lr_scheduler import CosineAnnealingLR
import time 
from tqdm import tqdm
import random
from tools import set_seeds, mkdir  
import pdb 
import pandas as pd

transformData = transformData()
io=dataIO() 
# set_seeds(42)



def save_model(G_net_model, save_dir, optimizer_G=None, ex=""):
    save_path=os.path.join(save_dir, "Model")
    mkdir(save_path)
    G_save_path = os.path.join(save_path,'Generator{}.pth'.format(ex))
    torch.save(G_net_model.cpu().state_dict(), G_save_path)
    G_net_model.cuda()

    if optimizer_G is not None:
        opt_G_save_path = os.path.join(save_path,'Optimizer_G{}.pth'.format(ex))
        torch.save(optimizer_G.state_dict(), opt_G_save_path)



def build_train_sampler(modality_list, data_root, batch_size, shuffle=True):
    dataset = Train_Data(root_dir = data_root, modality_list = modality_list)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=True, num_workers=4) 
    sampler = DataSampler(dataloader)
    print("data length: \n", dataset.length) 
    
    return sampler




total_iteration = 2e5
val_iteration = 1e3

batch_size = 8
eps=1e-8
lr=2e-4
psnr_max=0


data_root = "/home/data/zhiwen/dataset/All-in-One/" ### Path to place data
modality_list = ["PET", "CT", "MRI"] 
save_dir = "experiment/AMIR"



Generator = AMIR() 
Generator.cuda() 

train_sampler = build_train_sampler(modality_list, data_root, batch_size, shuffle=True) 
valid_loader = DataLoader(Test_Data(root_dir=data_root, use_num=32, modality_list=modality_list), batch_size=1, shuffle=False) 


optimizer_G = torch.optim.Adam(Generator.parameters(), lr=lr, betas=(0.9, 0.999), eps=1e-08) 
lr_scheduler_G = CosineAnnealingLR(optimizer_G, total_iteration, eta_min=1.0e-6)


L1 = nn.L1Loss().cuda() 


running_loss = []
eval_metrics={
    "psnr":[],
    "ssim":[], 
    "rmse":[]
    }


pbar = tqdm(total=int(total_iteration))

print("################ Train ################")
for iteration in list(range(1, int(total_iteration)+1)):


    l_G=[]
 

    in_pic,  label_pic, class_label = next(train_sampler)
    
    in_pic = in_pic.type(torch.FloatTensor).cuda()
    label_pic = label_pic.type(torch.FloatTensor).cuda() 


    Generator.train()
    optimizer_G.zero_grad() 

    restored, loss_load  = Generator(in_pic) 
    loss_l1 = L1(restored, label_pic) 
    loss_G = loss_l1 + 0.001*loss_load
    
    loss_G.backward()
    optimizer_G.step()

    l_G.append(loss_G.item())
    torch.cuda.empty_cache() 
    lr_scheduler_G.step()
    
    if iteration % val_iteration == 0: 
        psnr=0
        ssim=0
        rmse=0
        Generator.eval() 
        for counter,data in enumerate(tqdm(valid_loader)):
            v_in_pic, v_label_pic, modality, file_name = data 
            modality = modality[0] 
            file_name = file_name[0] 

            v_in_pic = v_in_pic.type(torch.FloatTensor).cuda() 
            v_label_pic = v_label_pic.type(torch.FloatTensor)

            with torch.no_grad():
                gen_img = Generator(v_in_pic) 
            
            gen_img = transformData.denormalize(gen_img, modality).detach().cpu() 
            
            v_label_pic = transformData.denormalize(v_label_pic, modality) 
            
            '''
            truncation for test image 
            CT:[-160, 240]
            '''
            
            gen_img = transformData.truncate_test(gen_img, modality) 
            v_label_pic = transformData.truncate_test(v_label_pic, modality) 
            
            data_range = v_label_pic.max()-v_label_pic.min()
            oneEval = compute_measure(gen_img, v_label_pic, data_range = data_range) 
            
            psnr+=oneEval[0]
            ssim+=oneEval[1] 
            rmse+=oneEval[2] 

            io.save(gen_img.clone().numpy().squeeze(), os.path.join(save_dir, "Gimg", "{}_{}.nii".format(file_name, modality) ))
            
            torch.cuda.empty_cache()
        c_psnr=psnr/(counter+1)
        c_ssim=ssim/(counter+1) 
        c_rmse=rmse/(counter+1) 
        
        eval_metrics['psnr'].append(c_psnr)
        eval_metrics['ssim'].append(c_ssim)  
        eval_metrics['rmse'].append(c_rmse) 
    
        save_model(G_net_model=Generator, save_dir=save_dir, optimizer_G=None, ex="_iteration_{}".format(iteration))
        if c_psnr>=psnr_max:
            psnr_max=c_psnr
            io.save("Best Iteration: {}, PSNR: {}, SSIM:{}, RMSE:{}".format(iteration, c_psnr, c_ssim, c_rmse),os.path.join(save_dir, "best.txt"))
            save_model(G_net_model=Generator, save_dir=save_dir, optimizer_G = optimizer_G, ex="_best")
        io.save(
            {'eval_metrics':eval_metrics},
            os.path.join(save_dir, "evaluationLoss.bin")
            ) 


    pbar.set_description("loss_G:{:6}, loss_load:{:6},  psnr:{:6}".format(loss_G.item(), loss_load.item(), eval_metrics['psnr'][-1] if len(eval_metrics['psnr'])>0 else 0)) 
    pbar.update()