0726-1429

diffusion.py

@@ -651,12 +651,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, 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")
 
     config = TrainConfig()
     config.device = f"cuda:{local_rank}"
@@ -668,11 +668,11 @@ def train(rank, world_size, local_world_size, master_addr, master_port):
         # config.world_size = world_size
         # print("ddpm21cm = DDPM21CM(config)")
         # print(f"config.device, torch.cuda.current_device() = {config.device}, {torch.cuda.current_device()}")
-    print("before dppm21cm")
+    #print("before dppm21cm")
     ddpm21cm = DDPM21CM(config)
     # print(f" num_image = {ddpm21cm.config.num_image} ".center(50, '-'))
     # print(f"run_name = {ddpm21cm.config.run_name}")
-    print(f"run_name {ddpm21cm.config.run_name}, global_rank {rank}, local_rank {local_rank}, current_device {torch.cuda.current_device()}, local_world_size {local_world_size}, world_size {world_size}")
+    #print(f"run_name {ddpm21cm.config.run_name}, global_rank {rank}, local_rank {local_rank}, current_device {torch.cuda.current_device()}, local_world_size {local_world_size}, world_size {world_size}")
     ddpm21cm.train()
     destroy_process_group()
 # %%

load_h5.py

@@ -168,7 +168,7 @@ class Dataset4h5(Dataset):
             param_start = time()
             params = f['params']['values'][idx]
             param_end = time()
-            print(f"ip_addr {socket.gethostbyname(socket.gethostname())}, cuda:{torch.cuda.current_device()}, CPU-pid {cpu_num}-{pid}: images {images.shape} & params {params.shape} loaded after {images_end-images_start:.3f}s & {param_end-param_start:.3f}s")
+            print(f"{socket.gethostbyname(socket.gethostname())}, cuda:{torch.cuda.current_device()}, CPU-pid {cpu_num}-{pid}: images {images.shape} & params {params.shape} loaded after {images_end-images_start:.3f}s & {param_end-param_start:.3f}s")
 
         return images, params
 

phoenix_diffusion.sbatch

@@ -3,6 +3,7 @@
 #SBATCH -A gts-jw254-coda20 
 #SBATCH -qembers
 #SBATCH -N2 --gpus-per-node=RTX_6000:3 # -C A100-80GB              # Number of nodes and cores per node required
+#SBATCH --ntasks-per-gpu=1
 #SBATCH --mem-per-gpu=32G                        # Memory per core
 #SBATCH -t 10:00                                    # Duration of the job (Ex: 15 mins)
 #SBATCH -oReport-%j                         # Combined output and error messages file