0728-1339

diffusion.py

@@ -81,10 +81,10 @@ def ddp_setup(rank: int, world_size: int, local_world_size, master_addr, master_
        world_size: Total number of processes
     """
 
-    print("inside ddp_setup")
+    #print("inside ddp_setup")
     os.environ["MASTER_ADDR"] = master_addr
     os.environ["MASTER_PORT"] = master_port
-    print("ddp_setup, rank =", rank)
+    #print("ddp_setup, rank =", rank)
     init_process_group(
             backend="nccl", 
             init_method=f"tcp://{master_addr}:{master_port}", 
@@ -574,7 +574,7 @@ class DDPM21CM:
                             'unet_state_dict': self.nn_model.module.state_dict(),
                             # 'ema_unet_state_dict': self.ema_model.state_dict(),
                             }
-                        save_name = self.config.save_name+f"-N{self.config.num_image}-device_count{self.config.world_size}-epoch{ep}"
+                        save_name = self.config.save_name+f"-N{self.config.num_image}-device_count{self.config.world_size}-epoch{ep}-{socket.gethostbyname(socket.gethostname())}"
                         torch.save(model_state, save_name)
                         print(f'cuda:{torch.cuda.current_device()} saved model at ' + save_name)
                         # print('saved model at ' + config.save_dir + f"model_epoch_{ep}_test_{config.run_name}.pth")
@@ -658,12 +658,12 @@ class DDPM21CM:
 
 #num_train_image_list = [6000]#[60]#[8000]#[1000]#[100]#
 def train(rank, world_size, local_world_size, master_addr, master_port):
-    print("before ddp_setup")
+    #print("before ddp_setup")
     ddp_setup(rank, world_size, local_world_size, master_addr, master_port)
-    print("after ddp_setup")
+    #print("after ddp_setup")
     local_rank = rank % local_world_size
     torch.cuda.set_device(local_rank)
-    print("after set device")
+    #print("after set device")
     print(f"rank = {rank}, local_rank = {local_rank}, world_size = {world_size}, local_world_size = {local_world_size}")
 
     config = TrainConfig()

load_h5.py

@@ -96,10 +96,10 @@ class Dataset4h5(Dataset):
             print(f"dataset content: {f.keys()}")
             max_num_image = len(f['brightness_temp'])#.shape[0]
             field_shape = f['brightness_temp'].shape[1:]
-            print(f"{max_num_image} images of shape {field_shape} can be loaded")
             #print(f"field.shape = {field_shape}")
             self.params_keys = list(f['params']['keys'])
-            print(f"params keys = {self.params_keys}")
+            print(f"{max_num_image} images of shape {field_shape} can be loaded with different params.keys {self.params_keys}")
+            #print(f"params keys = {self.params_keys}")
 
         # if self.idx is None:
         #     if self.shuffle: