0803-182144

diffusion.py

@@ -260,7 +260,7 @@ class TrainConfig:
     stride = (2,4) if dim == 2 else (2,2,2)
     num_image = 3000#300#3000#6000#30#60#6000#1000#2000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560
     batch_size = 10#50#10#50#20#50#1#2#50#20#2#100 # 10
-    n_epoch = 50#5#50#5#50#100#50#100#30#120#5#4# 10#50#20#20#2#5#25 # 120
+    n_epoch = 1#50#5#50#5#50#100#50#100#30#120#5#4# 10#50#20#20#2#5#25 # 120
     HII_DIM = 64
     num_redshift = 512#64#512#64#512#64#256CUDAoom#128#64#512#128#64#512#256#256#64#512#128
     channel = 1
@@ -711,7 +711,7 @@ def generate_samples(rank, world_size, local_world_size, master_addr, master_por
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--train", type=int, required=False, help="whether to train the model", default=1)
-    parser.add_argument("--sample", type=int, required=False, help="whether to sample", default=0)
+    #parser.add_argument("--sample", type=int, required=False, help="whether to sample", default=0)
     parser.add_argument("--resume", type=str, required=False, help="filename of the model to resume", default=False)
     parser.add_argument("--num_new_img_per_gpu", type=int, required=False, default=4)
     parser.add_argument("--max_num_img_per_gpu", type=int, required=False, default=2)
@@ -725,7 +725,7 @@ if __name__ == "__main__":
     world_size = local_world_size * total_nodes #6#int(os.environ["SLURM_NTASKS"])
 
     ############################ training ################################
-    if args.train:
+    if args.train == 1:
         print(f" training, ip_addr = {socket.gethostbyname(socket.gethostname())}, master_addr = {master_addr}, local_world_size = {local_world_size}, world_size = {world_size} ".center(120,'-'))
         mp.spawn(
                 train, 
@@ -734,7 +734,7 @@ if __name__ == "__main__":
                 join=True,
                 )
     ############################ sampling ################################
-    if args.sample:
+    if args.train == 0:
         num_new_img_per_gpu = args.num_new_img_per_gpu#200#4#200
         max_num_img_per_gpu = args.max_num_img_per_gpu#40#2#20
         config = TrainConfig()
@@ -748,10 +748,10 @@ if __name__ == "__main__":
         # return_dict = manager.dict()
         params_pairs = [
             (4.4, 131.341),
-            #(5.6, 19.037),
-            #(4.699, 30),
-            #(5.477, 200),
-            #(4.8, 131.341),
+            (5.6, 19.037),
+            (4.699, 30),
+            (5.477, 200),
+            (4.8, 131.341),
         ]
 
         for params in params_pairs:

phoenix_diffusion.sbatch

@@ -2,10 +2,10 @@
 #SBATCH -J diffusion # Job name
 #SBATCH -A gts-jw254-coda20 
 #SBATCH -qembers
-#SBATCH -N1 --gpus-per-node=RTX_6000:2 # -C A100-80GB              # Number of nodes and cores per node required
+#SBATCH -N8 --gpus-per-node=RTX_6000:1 # -C A100-80GB              # Number of nodes and cores per node required
 #SBATCH --ntasks-per-node=1
-#SBATCH --mem-per-gpu=16G                        # Memory per core
-#SBATCH -t 03:00:00                                    # Duration of the job (Ex: 15 mins)
+#SBATCH --mem-per-gpu=8G                        # Memory per core
+#SBATCH -t 00:05:00                                    # Duration of the job (Ex: 15 mins)
 #SBATCH -oReport-%j                         # Combined output and error messages file
 #SBATCH --error=error-%j
 #SBATCH --mail-type=BEGIN,END,FAIL              # Mail preferences
@@ -30,10 +30,9 @@ export MASTER_PORT=$MASTER_PORT
 
 srun python diffusion.py \
     --train 1 \
-    --sample 0 \
-    --resume outputs/model_state-N3000-device_count3-node2-epoch49-172.27.145.67 \
-    --num_new_img_per_gpu 200 \
-    --max_num_img_per_gpu 40 \
+    --resume outputs/model_state-N3000-device_count1-node4-epoch49-172.27.149.191 \
+    --num_new_img_per_gpu 50 \
+    --max_num_img_per_gpu 10 \
 
 ######################################################################################