0710-1816

diffusion.py

@@ -361,7 +361,7 @@ class DDPM21CM:
         dataset = Dataset4h5(self.config.dataset_name, num_image=self.config.num_image, HII_DIM=self.config.HII_DIM, num_redshift=self.config.num_redshift, drop_prob=self.config.drop_prob, dim=self.config.dim, ranges_dict=self.ranges_dict)
         # self.shape_loaded = dataset.images.shape
         # print("shape_loaded =", self.shape_loaded)
-        # print(f"load, current_device = {torch.cuda.current_device()}")
+        # print(f"load, current_device() = {torch.cuda.current_device()}")
         self.dataloader = DataLoader(
             dataset=dataset, 
             batch_size=self.config.batch_size, 
@@ -394,7 +394,8 @@ class DDPM21CM:
             # distributed_type="MULTI_GPU",
         )
         # print("!!!!!!!!!!!!!!!!!!!self.accelerator.device:", self.accelerator.device)
-        if self.accelerator.is_main_process:
+        # if self.accelerator.is_main_process:
+        if torch.cuda.current_device() == 0:
             if self.config.output_dir is not None:
                 os.makedirs(self.config.output_dir, exist_ok=True)
             if self.config.push_to_hub:
@@ -414,7 +415,7 @@ class DDPM21CM:
             self.nn_model, self.optimizer, self.lr_scheduler
             )
 
-        print("!!!!!!!!!!!!!!!!, after prepare, self.dataloader.sampler =", self.dataloader.sampler)
+        # print("!!!!!!!!!!!!!!!!, after prepare, self.dataloader.sampler =", self.dataloader.sampler)
         # print("!!!!!!!!!!!!!!!!, after prepare, self.dataloader.batch_sampler =", self.dataloader.batch_sampler)
         # print("!!!!!!!!!!!!!!!!, after prepare, self.dataloader.DistributedSampler =", self.dataloader.DistributedSampler)
 
@@ -468,7 +469,8 @@ class DDPM21CM:
 
     def save(self, ep):
         # save model
-        if self.accelerator.is_main_process:
+        # if self.accelerator.is_main_process:
+        if torch.cuda.current_device() == 0:
             if ep == self.config.n_epoch-1 or (ep+1)*self.config.save_freq==1:
                 self.nn_model.eval()
                 with torch.no_grad():
@@ -486,7 +488,7 @@ class DDPM21CM:
                             # 'ema_unet_state_dict': self.ema_model.state_dict(),
                             }
                         torch.save(model_state, self.config.save_name+f"-N{self.config.num_image}")
-                        print('saved model at ' + self.config.save_name+f"-N{self.config.num_image}")
+                        print(f'device {torch.cuda.current_device()} saved model at ' + self.config.save_name+f"-N{self.config.num_image}")
                         # print('saved model at ' + config.save_dir + f"model_epoch_{ep}_test_{config.run_name}.pth")
 
     # def rescale(self, value, type='params', to_ranges=[0,1]):