0721-2348

diffusion.py

@@ -341,6 +341,8 @@ class DDPM21CM:
             # print(f"resumed nn_model from {config.resume}")
             self.nn_model.module.load_state_dict(torch.load(config.resume)['unet_state_dict'])
             print(f"device {torch.cuda.current_device()} resumed nn_model from {config.resume}")
+        else:
+            print(f"device {torch.cuda.current_device()} initialized nn_model randomly")
 
         self.number_of_params = sum(x.numel() for x in self.nn_model.parameters())
         print(f" Number of parameters for nn_model: {self.number_of_params} ".center(100,'-'))
@@ -614,7 +616,7 @@ def train(rank, world_size):
 
         
 if __name__ == "__main__":
-    world_size = torch.cuda.device_count()
+    world_size = 2#torch.cuda.device_count()
     print(f" training, world_size = {world_size} ".center(100,'-'))
     # torch.multiprocessing.set_start_method("spawn")
     # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)
@@ -691,7 +693,7 @@ if __name__ == "__main__":
 
     for num_image in num_train_image_list:
         config.num_image = num_image // world_size
-        config.resume = f"./outputs/model_state-N8000-device_count4-epoch{config.n_epoch-1}"
+        config.resume = f"./outputs/model_state-N4000-device_count2-epoch{config.n_epoch-1}"
 
         # print("ddpm21cm = DDPM21CM(config)")
         manager = mp.Manager()