modify gpu allocation

Author: mingukkang

src/loader.py

@@ -42,15 +42,15 @@ def prepare_train_eval(local_rank, gpus_per_node, world_size, run_name, train_co
     prev_ada_p, step, best_step, best_fid, best_fid_checkpoint_path, mu, sigma, inception_model = None, 0, 0, None, None, None, None, None
 
     if cfgs.distributed_data_parallel:
-        rank = cfgs.nr*(gpus_per_node) + local_rank
-        print("Use GPU: {} for training.".format(rank))
-        setup(rank, world_size)
-        torch.cuda.set_device(rank)
+        global_rank = cfgs.nr*(gpus_per_node) + local_rank
+        print("Use GPU: {} for training.".format(global_rank))
+        setup(global_rank, world_size)
+        torch.cuda.set_device(local_rank)
     else:
-        rank = local_rank
+        global_rank = local_rank
 
-    writer = SummaryWriter(log_dir=join('./logs', run_name)) if rank == 0 else None
-    if rank == 0:
+    writer = SummaryWriter(log_dir=join('./logs', run_name)) if global_rank == 0 else None
+    if global_rank == 0:
         logger = make_logger(run_name, None)
         logger.info('Run name : {run_name}'.format(run_name=run_name))
         logger.info(train_config)
@@ -59,19 +59,19 @@ def prepare_train_eval(local_rank, gpus_per_node, world_size, run_name, train_co
         logger = None
 
     ##### load dataset #####
-    if rank == 0: logger.info('Load train datasets...')
+    if global_rank == 0: logger.info('Load train datasets...')
     train_dataset = LoadDataset(cfgs.dataset_name, cfgs.data_path, train=True, download=True, resize_size=cfgs.img_size,
                                 hdf5_path=hdf5_path_train, random_flip=cfgs.random_flip_preprocessing)
     if cfgs.reduce_train_dataset < 1.0:
         num_train = int(cfgs.reduce_train_dataset*len(train_dataset))
         train_dataset, _ = torch.utils.data.random_split(train_dataset, [num_train, len(train_dataset) - num_train])
-    if rank == 0: logger.info('Train dataset size : {dataset_size}'.format(dataset_size=len(train_dataset)))
+    if global_rank == 0: logger.info('Train dataset size : {dataset_size}'.format(dataset_size=len(train_dataset)))
 
-    if rank == 0: logger.info('Load {mode} datasets...'.format(mode=cfgs.eval_type))
+    if global_rank == 0: logger.info('Load {mode} datasets...'.format(mode=cfgs.eval_type))
     eval_mode = True if cfgs.eval_type == 'train' else False
     eval_dataset = LoadDataset(cfgs.dataset_name, cfgs.data_path, train=eval_mode, download=True, resize_size=cfgs.img_size,
                                hdf5_path=None, random_flip=False)
-    if rank == 0: logger.info('Eval dataset size : {dataset_size}'.format(dataset_size=len(eval_dataset)))
+    if global_rank == 0: logger.info('Eval dataset size : {dataset_size}'.format(dataset_size=len(eval_dataset)))
 
     if cfgs.distributed_data_parallel:
         train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
@@ -84,31 +84,31 @@ def prepare_train_eval(local_rank, gpus_per_node, world_size, run_name, train_co
     eval_dataloader = DataLoader(eval_dataset, batch_size=cfgs.batch_size, shuffle=False, pin_memory=True, num_workers=cfgs.num_workers, drop_last=False)
 
     ##### build model #####
-    if rank == 0: logger.info('Build model...')
+    if global_rank == 0: logger.info('Build model...')
     module = __import__('models.{architecture}'.format(architecture=cfgs.architecture), fromlist=['something'])
-    if rank == 0: logger.info('Modules are located on models.{architecture}.'.format(architecture=cfgs.architecture))
+    if global_rank == 0: logger.info('Modules are located on models.{architecture}.'.format(architecture=cfgs.architecture))
     Gen = module.Generator(cfgs.z_dim, cfgs.shared_dim, cfgs.img_size, cfgs.g_conv_dim, cfgs.g_spectral_norm, cfgs.attention,
                            cfgs.attention_after_nth_gen_block, cfgs.activation_fn, cfgs.conditional_strategy, cfgs.num_classes,
-                           cfgs.g_init, cfgs.G_depth, cfgs.mixed_precision).to(rank)
+                           cfgs.g_init, cfgs.G_depth, cfgs.mixed_precision).to(local_rank)
 
     Dis = module.Discriminator(cfgs.img_size, cfgs.d_conv_dim, cfgs.d_spectral_norm, cfgs.attention, cfgs.attention_after_nth_dis_block,
                                cfgs.activation_fn, cfgs.conditional_strategy, cfgs.hypersphere_dim, cfgs.num_classes, cfgs.nonlinear_embed,
-                               cfgs.normalize_embed, cfgs.d_init, cfgs.D_depth, cfgs.mixed_precision).to(rank)
+                               cfgs.normalize_embed, cfgs.d_init, cfgs.D_depth, cfgs.mixed_precision).to(local_rank)
 
     if cfgs.ema:
-        if rank == 0: logger.info('Prepare EMA for G with decay of {}.'.format(cfgs.ema_decay))
+        if global_rank == 0: logger.info('Prepare EMA for G with decay of {}.'.format(cfgs.ema_decay))
         Gen_copy = module.Generator(cfgs.z_dim, cfgs.shared_dim, cfgs.img_size, cfgs.g_conv_dim, cfgs.g_spectral_norm, cfgs.attention,
                                     cfgs.attention_after_nth_gen_block, cfgs.activation_fn, cfgs.conditional_strategy, cfgs.num_classes,
-                                    initialize=False, G_depth=cfgs.G_depth, mixed_precision=cfgs.mixed_precision).to(rank)
+                                    initialize=False, G_depth=cfgs.G_depth, mixed_precision=cfgs.mixed_precision).to(local_rank)
         Gen_ema = ema(Gen, Gen_copy, cfgs.ema_decay, cfgs.ema_start)
     else:
         Gen_copy, Gen_ema = None, None
 
-    if rank == 0: logger.info(count_parameters(Gen))
-    if rank == 0: logger.info(Gen)
+    if global_rank == 0: logger.info(count_parameters(Gen))
+    if global_rank == 0: logger.info(Gen)
 
-    if rank == 0: logger.info(count_parameters(Dis))
-    if rank == 0: logger.info(Dis)
+    if global_rank == 0: logger.info(count_parameters(Dis))
+    if global_rank == 0: logger.info(Dis)
 
 
     ### define loss functions and optimizers
@@ -144,18 +144,18 @@ def prepare_train_eval(local_rank, gpus_per_node, world_size, run_name, train_co
         Gen, G_optimizer, trained_seed, run_name, step, prev_ada_p = load_checkpoint(Gen, G_optimizer, g_checkpoint_dir)
         Dis, D_optimizer, trained_seed, run_name, step, prev_ada_p, best_step, best_fid, best_fid_checkpoint_path =\
             load_checkpoint(Dis, D_optimizer, d_checkpoint_dir, metric=True)
-        if rank == 0: logger = make_logger(run_name, None)
+        if global_rank == 0: logger = make_logger(run_name, None)
         if cfgs.ema:
             g_ema_checkpoint_dir = glob.glob(join(checkpoint_dir, "model=G_ema-{when}-weights-step*.pth".format(when=when)))[0]
             Gen_copy = load_checkpoint(Gen_copy, None, g_ema_checkpoint_dir, ema=True)
             Gen_ema.source, Gen_ema.target = Gen, Gen_copy
 
-        writer = SummaryWriter(log_dir=join('./logs', run_name)) if rank ==0 else None
+        writer = SummaryWriter(log_dir=join('./logs', run_name)) if global_rank == 0 else None
         if cfgs.train_configs['train']:
             assert cfgs.seed == trained_seed, "Seed for sampling random numbers should be same!"
 
-        if rank == 0: logger.info('Generator checkpoint is {}'.format(g_checkpoint_dir))
-        if rank == 0: logger.info('Discriminator checkpoint is {}'.format(d_checkpoint_dir))
+        if global_rank == 0: logger.info('Generator checkpoint is {}'.format(g_checkpoint_dir))
+        if global_rank == 0: logger.info('Discriminator checkpoint is {}'.format(d_checkpoint_dir))
         if cfgs.freeze_layers > -1 :
             prev_ada_p, step, best_step, best_fid, best_fid_checkpoint_path = None, 0, 0, None, None
 
@@ -170,30 +170,30 @@ def prepare_train_eval(local_rank, gpus_per_node, world_size, run_name, train_co
                 if cfgs.ema:
                     Gen_copy = torch.nn.SyncBatchNorm.convert_sync_batchnorm(Gen_copy, process_group)
 
-            Gen = DDP(Gen, device_ids=[rank])
-            Dis = DDP(Dis, device_ids=[rank])
+            Gen = DDP(Gen, device_ids=[local_rank])
+            Dis = DDP(Dis, device_ids=[local_rank])
             if cfgs.ema:
-                Gen_copy = DDP(Gen_copy, device_ids=[rank])
+                Gen_copy = DDP(Gen_copy, device_ids=[local_rank])
         else:
-            Gen = DataParallel(Gen, output_device=rank)
-            Dis = DataParallel(Dis, output_device=rank)
+            Gen = DataParallel(Gen, output_device=local_rank)
+            Dis = DataParallel(Dis, output_device=local_rank)
             if cfgs.ema:
-                Gen_copy = DataParallel(Gen_copy, output_device=rank)
+                Gen_copy = DataParallel(Gen_copy, output_device=local_rank)
 
             if cfgs.synchronized_bn:
-                Gen = convert_model(Gen).to(rank)
-                Dis = convert_model(Dis).to(rank)
+                Gen = convert_model(Gen).to(local_rank)
+                Dis = convert_model(Dis).to(local_rank)
                 if cfgs.ema:
-                    Gen_copy = convert_model(Gen_copy).to(rank)
+                    Gen_copy = convert_model(Gen_copy).to(local_rank)
 
     ##### load the inception network and prepare first/secend moments for calculating FID #####
     if cfgs.eval:
-        inception_model = InceptionV3().to(rank)
+        inception_model = InceptionV3().to(local_rank)
         if world_size > 1 and cfgs.distributed_data_parallel:
             toggle_grad(inception_model, on=True)
-            inception_model = DDP(inception_model, device_ids=[rank], broadcast_buffers=False, find_unused_parameters=True)
+            inception_model = DDP(inception_model, device_ids=[local_rank], broadcast_buffers=False, find_unused_parameters=True)
         elif world_size > 1 and cfgs.distributed_data_parallel is False:
-            inception_model = DataParallel(inception_model, output_device=rank)
+            inception_model = DataParallel(inception_model, output_device=local_rank)
         else:
             pass
 
@@ -204,7 +204,7 @@ def prepare_train_eval(local_rank, gpus_per_node, world_size, run_name, train_co
                                               splits=1,
                                               run_name=run_name,
                                               logger=logger,
-                                              device=rank)
+                                              device=local_rank)
 
     worker = make_worker(
         cfgs=cfgs,
@@ -227,7 +227,8 @@ def prepare_train_eval(local_rank, gpus_per_node, world_size, run_name, train_co
         G_loss=G_loss[cfgs.adv_loss],
         D_loss=D_loss[cfgs.adv_loss],
         prev_ada_p=prev_ada_p,
-        rank=rank,
+        global_rank=global_rank,
+        local_rank=local_rank,
         bn_stat_OnTheFly=cfgs.bn_stat_OnTheFly,
         checkpoint_dir=checkpoint_dir,
         mu=mu,