0719-1039

diffusion.py

@@ -236,12 +236,12 @@ class TrainConfig:
 
     # dim = 2
     dim = 2
-    stride = (2,4) if dim == 2 else (2,2,4)
+    stride = (2,2) if dim == 2 else (2,2,4)
     num_image = 1000#2000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560
     batch_size = 10#1#2#50#20#2#100 # 10
-    n_epoch = 10#5#4# 10#50#20#20#2#5#25 # 120
+    n_epoch = 30#5#4# 10#50#20#20#2#5#25 # 120
     HII_DIM = 64
-    num_redshift = 512#128#64#512#256#256#64#512#128
+    num_redshift = 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)
 
@@ -581,7 +581,7 @@ class DDPM21CM:
         return x_last
 # %%
 
-num_train_image_list = [5000]
+num_train_image_list = [4000]
 
 def train(rank, world_size):
     config = TrainConfig()