0804-202215

context_unet.py

@@ -330,7 +330,7 @@ class ContextUnet(nn.Module):
             elif image_size == 128:
                 channel_mult = (1, 1, 2, 3, 4)
             elif image_size == 64:
-                channel_mult = (2,2,4,4,4)#(1, 2, 4)#(2,4,4,4,8)#(1, 2, 2, 4, 4)#(1, 2, 2, 4, 8)#(1, 1, 2, 2, 4, 4)#(1, 2, 4, 8, 16)#(1, 2, 3, 4)#(1, 2, 4, 6, 8)#(1, 2, 2, 4)#(1, 2, 8, 8, 8)#(1, 2, 4)#(1, 2, 2, 4)#(0.5,1,2,2,4,4)#(1, 1, 2, 2, 4, 4)#
+                channel_mult = (0.5,1,1,2,2)#(1,1,2)#(1,2)#(1,1,2,2)#(1,1,2,2,4)#(2,2,4,4,4)#(1, 2, 4)#(2,4,4,4,8)#(1, 2, 2, 4, 4)#(1, 2, 2, 4, 8)#(1, 1, 2, 2, 4, 4)#(1, 2, 4, 8, 16)#(1, 2, 3, 4)#(1, 2, 4, 6, 8)#(1, 2, 2, 4)#(1, 2, 8, 8, 8)#(1, 2, 4)#(1, 2, 2, 4)#(0.5,1,2,2,4,4)#(1, 1, 2, 2, 4, 4)#
             elif image_size == 32:
                 channel_mult = (1, 2, 2, 4)
             elif image_size == 28:
@@ -550,4 +550,4 @@ class ContextUnet(nn.Module):
         h = self.out(h)
         # print("self.out(h)", "h.shape =", h.shape)
 
-        return h 
+        return h 

diffusion.py

@@ -256,13 +256,13 @@ class TrainConfig:
     # repeat = 2
 
     # dim = 2
-    dim = 2
-    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#1#50#5#50#5#50#100#50#100#30#120#5#4# 10#50#20#20#2#5#25 # 120
+    dim = 3#2
+    stride = (2,4) if dim == 2 else (2,2,4)
+    num_image = 480#1200#120#3000#300#3000#6000#30#60#6000#1000#2000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560
+    batch_size = 1#10#50#10#50#20#50#1#2#50#20#2#100 # 10
+    n_epoch = 50#1#50#10#1#50#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
+    num_redshift = 512#256#512#256#512#256#512#64#512#64#512#64#256CUDAoom#128#64#512#128#64#512#256#256#64#512#128
     channel = 1
     img_shape = (channel, HII_DIM, num_redshift) if dim == 2 else (channel, HII_DIM, HII_DIM, num_redshift)
 
@@ -303,7 +303,7 @@ class TrainConfig:
     # params =  params
     # data_dir = './data' # data directory
 
-    use_fp16 = False
+    use_fp16 = True 
     dtype = torch.float16 if use_fp16 else torch.float32
     mixed_precision = "fp16"
     gradient_accumulation_steps = 1

phoenix_diffusion.sbatch

@@ -2,10 +2,10 @@
 #SBATCH -J diffusion # Job name
 #SBATCH -A gts-jw254-coda20 
 #SBATCH -qembers
-#SBATCH -N8 --gpus-per-node=RTX_6000:1 # -C A100-80GB              # Number of nodes and cores per node required
+#SBATCH -N10 --gpus-per-node=V100:1 -C V100-32GB              # Number of nodes and cores per node required
 #SBATCH --ntasks-per-node=1
-#SBATCH --mem-per-gpu=8G                        # Memory per core
-#SBATCH -t 03:10:00                                    # Duration of the job (Ex: 15 mins)
+#SBATCH --mem-per-gpu=16G                        # Memory per core
+#SBATCH -t 08:00: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,7 +30,7 @@ export MASTER_PORT=$MASTER_PORT
 
 srun python diffusion.py \
     --train 1 \
-    --resume outputs/model_state-N3000-device_count1-node4-epoch49-172.27.149.191 \
+    --resume outputs/model_state-N3000-device_count1-node8-epoch49-172.27.149.181 \
     --num_new_img_per_gpu 50 \
     --max_num_img_per_gpu 10 \