import os import torch import torchvision from torch import nn from torchvision.utils import save_image from torch.utils.data import DataLoader from torch.optim import Adam, SGD from Diffusion.diffuser import DeformDDPM from Diffusion.networks import get_net, STN from torchvision.transforms import Lambda import Diffusion.losses as losses import random import glob import numpy as np import utils from Dataloader.dataloader0 import get_dataloader from Dataloader.dataloader_utils import thresh_img import yaml import argparse #################### import torch.multiprocessing as mp from torch.utils.data.distributed import DistributedSampler from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed import init_process_group, destroy_process_group ############### def ddp_setup(rank, world_size): """ Args: rank: Unique identifier of each process world_size: Total number of processes """ os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "12355" init_process_group(backend="nccl", rank=rank, world_size=world_size) torch.cuda.set_device(rank) use_parallel=False use_distributed = False EPS = 1e-5 parser = argparse.ArgumentParser() # config_file_path = 'Config/config_cmr.yaml' parser.add_argument( "--config", "-C", help="Path for the config file", type=str, default="Config/config_cmr.yaml", # default="Config/config_lct.yaml", required=False, ) args = parser.parse_args() #======================================================================================================================= def main_train(rank,world_size): ddp_setup(rank,world_size) gpu_id = rank # Load the YAML file into a dictionary with open(args.config, 'r') as file: hyp_parameters = yaml.safe_load(file) print(hyp_parameters) # epoch_per_save=10 epoch_per_save=hyp_parameters['epoch_per_save'] data_name=hyp_parameters['data_name'] net_name = hyp_parameters['net_name'] Net=get_net(net_name) suffix_pth=f'_{data_name}_{net_name}.pth' model_save_path = os.path.join('Models',f'{data_name}_{net_name}/') model_dir=model_save_path transformer=utils.get_transformer(img_sz=hyp_parameters["ndims"]*[hyp_parameters['img_size']]) Data_Loader=get_dataloader(data_name=hyp_parameters['data_name'], mode='train') tsfm = torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), ]) dataset = Data_Loader(target_res = [hyp_parameters["img_size"]]*hyp_parameters["ndims"], transforms=None, noise_scale=hyp_parameters['noise_scale']) train_loader = DataLoader( dataset, batch_size=hyp_parameters['batchsize'], # shuffle=False, shuffle=True, drop_last=True, ) Deformddpm = DeformDDPM( network=Net(n_steps=hyp_parameters["timesteps"], ndims=hyp_parameters["ndims"], num_input_chn=1), n_steps=hyp_parameters["timesteps"], image_chw=[1] + [hyp_parameters["img_size"]]*hyp_parameters["ndims"], device=hyp_parameters["device"], batch_size=hyp_parameters["batchsize"], img_pad_mode=hyp_parameters["img_pad_mode"], v_scale=hyp_parameters["v_scale"], ) ddf_stn = STN( img_sz=hyp_parameters["img_size"], ndims=hyp_parameters["ndims"], # padding_mode="zeros", padding_mode=hyp_parameters["padding_mode"], device=hyp_parameters["device"], ) # Deformddpm.to(hyp_parameters["device"]) # ddf_stn.to(hyp_parameters["device"]) # if use_distributed: # torch.distributed.init_process_group(backend='nccl') # Deformddpm = nn.parallel.DistributedDataParallel(Deformddpm, device_ids=[torch.cuda.current_device()]) # ddf_stn = nn.parallel.DistributedDataParallel(ddf_stn, device_ids=[torch.cuda.current_device()]) # elif use_parallel: # Deformddpm = nn.DataParallel(Deformddpm) # ddf_stn = nn.DataParallel(ddf_stn) Deformddpm.to(rank) Deformddpm = DDP(Deformddpm, device_ids=[rank]) ddf_stn.to(rank) # ddf_stn = DDP(ddf_stn, device_ids=[rank]) # mse = nn.MSELoss() loss_reg = losses.Grad(penalty=['l1', 'negdetj'], ndims=hyp_parameters["ndims"]) loss_dist = losses.MRSE(img_sz=hyp_parameters["img_size"]) # loss_ang = losses.MRSE(img_sz=hyp_parameters["img_size"]) loss_ang = losses.NCC(img_sz=hyp_parameters["img_size"]) optimizer = Adam(Deformddpm.parameters(), lr=hyp_parameters["lr"]) # hyp_parameters["lr"]=0.00000001 # # optimizer = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"], momentum=0.95) # optimizer = SGD(Deformddpm.parameters(), lr=hyp_parameters["lr"], momentum=0.9) # # LR scheduler ----- YHM # scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, hyp_parameters["lr"], hyp_parameters["lr"]*10, step_size_up=500, step_size_down=500, mode='triangular', gamma=1.0, scale_fn=None, scale_mode='cycle', cycle_momentum=True, base_momentum=0.8, max_momentum=0.9, last_epoch=-1) # Deformddpm.network.load_state_dict(torch.load('/home/data/jzheng/Adaptive_Motion_Generator-master/models/1000.pth')) # check for existing models if not os.path.exists(model_dir): os.makedirs(model_dir, exist_ok=True) model_files = glob.glob(os.path.join(model_dir, "*.pth")) model_files.sort() print(model_files) if model_files: # if there are any model files, load the most recent one latest_model_file = model_files[-1] # Deformddpm.network.load_state_dict(torch.load(latest_model_file)) if use_parallel: Deformddpm.module.load_state_dict(torch.load(latest_model_file), strict=False) else: Deformddpm.load_state_dict(torch.load(latest_model_file), strict=False) # get the epoch number from the filename and add 1 to set as initial_epoch initial_epoch = int(os.path.basename(latest_model_file).split('.')[0][:6]) + 1 else: initial_epoch = 0 print('len_train_data: ',len(dataset)) for epoch in range(initial_epoch,hyp_parameters["epoch"]): epoch_loss_tot = 0.0 epoch_loss_gen_d = 0.0 epoch_loss_gen_a = 0.0 epoch_loss_reg = 0.0 # Set model inside to train model Deformddpm.train() for step, batch in enumerate(train_loader): # x0, _ = batch x0, _, _ = batch x0 = x0.to(hyp_parameters["device"]).type(torch.float32) n = x0.size()[0] # batch_size -> n x0 = x0.to(hyp_parameters["device"]) # random deformation + rotation if hyp_parameters["ndims"]>2: if np.random.uniform(0,1)<0.6: x0 = utils.random_resample(x0, deform_scale=0) x0 = transformer(x0) if hyp_parameters['noise_scale']>0: x0 = thresh_img(x0, [0, 2*hyp_parameters['noise_scale']]) x0 = x0 * (np.random.normal(1, hyp_parameters['noise_scale'] * 1)) + np.random.normal(0, hyp_parameters['noise_scale'] * 1) # Picking some noise for each of the images in the batch, a timestep and the respective alpha_bars t = torch.randint(0, hyp_parameters["timesteps"], (n,)).to( hyp_parameters["device"] ) # pick up a seq of rand number from 0 to 'timestep' if use_parallel: # # noisy_imgs, dvf_I = ddf_enc(img= x0, t) # noisy_imgs, dvf_I,_ = Deformddpm.module.diffuse(x0, t) # pre_dvf_I = Deformddpm.backward(noisy_imgs, t.reshape(16, -1)) pre_dvf_I, _ = Deformddpm.module(x0, t) else: # # noisy_imgs, dvf_I = ddf_enc(img= x0, t) # noisy_imgs, dvf_I,_ = Deformddpm.diffuse(x0, t) # pre_dvf_I = Deformddpm.backward(noisy_imgs, t.reshape(16, -1)) pre_dvf_I,dvf_I = Deformddpm(x0, t) loss_tot=0 loss_ddf = loss_reg(pre_dvf_I) trm_pred = ddf_stn(pre_dvf_I, dvf_I) loss_gen_d = loss_dist(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None) loss_gen_a = loss_ang(pred=trm_pred,inv_lab=dvf_I,ddf_stn=None) loss_tot += 1.0 * loss_gen_d + 1.0 * loss_gen_a loss_tot +=10 * loss_ddf optimizer.zero_grad() loss_tot.backward() optimizer.step() epoch_loss_tot += loss_tot.item() * len(x0) / len(train_loader.dataset) epoch_loss_gen_d += loss_gen_d.item() * len(x0) / len(train_loader.dataset) epoch_loss_gen_a += loss_gen_a.item() * len(x0) / len(train_loader.dataset) epoch_loss_reg += loss_ddf.item() * len(x0) / len(train_loader.dataset) # print('step:',step,':', loss_tot.item(),'=',loss_gen_a.item(),'+', loss_gen_d.item(),'+',loss_ddf.item()) if gpu_id == 0: print(epoch,':', epoch_loss_tot,'=',epoch_loss_gen_a,'+', epoch_loss_gen_d,'+',epoch_loss_reg, ' (ang+dist+regul)') # # LR schedular step ----- YHM # scheduler.step() if 0 == epoch % epoch_per_save: save_dir=model_save_path + str(epoch).rjust(6, '0') + suffix_pth if os.path.exists(model_save_path): print(f"saved in {save_dir}") else: os.makedirs(os.path.dirname(model_save_path)) # break # FOR TESTING if use_parallel: torch.save(Deformddpm.module.state_dict(), save_dir) elif gpu_id == 0: torch.save(Deformddpm.module.state_dict(), save_dir) if __name__ == "__main__": world_size = torch.cuda.device_count() print(f"world size = {world_size}") mp.spawn(main_train,args = (world_size,),nprocs = world_size)