0710-2338

diffusion.py

@@ -230,7 +230,7 @@ class TrainConfig:
     hub_private_repo = False
     dataset_name = "/storage/home/hcoda1/3/bxia34/scratch/LEN128-DIM64-CUB8.h5"
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    # world_size = torch.cuda.device_count()
+    world_size = torch.cuda.device_count()
     # repeat = 2
 
     # dim = 2
@@ -351,7 +351,7 @@ class DDPM21CM:
         self.lr_scheduler = get_cosine_schedule_with_warmup(
             optimizer=self.optimizer,
             num_warmup_steps=config.lr_warmup_steps,
-            num_training_steps=(int(config.num_image/config.batch_size) * config.n_epoch),
+            num_training_steps=int(config.num_image / config.world_size / config.batch_size * config.n_epoch),
             # num_training_steps=(len(self.dataloader) * config.n_epoch),
         )
 
@@ -558,6 +558,8 @@ class DDPM21CM:
 # %%
 def main(rank, world_size):
     config = TrainConfig()
+    config.world_size = world_size
+
     ddp_setup(rank, world_size)
     
     num_image_list = [5000]#[200]#[1600,3200,6400,12800,25600]