0716-2205

diffusion.py

@@ -236,12 +236,12 @@ class TrainConfig:
 
     # dim = 2
     dim = 3
-    stride = (2,2) if dim == 2 else (2,2,2)
+    stride = (2,2) if dim == 2 else (2,2,4)
     num_image = 1000#32000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560
     batch_size = 1#2#50#20#2#100 # 10
     n_epoch = 8#4# 10#50#20#20#2#5#25 # 120
     HII_DIM = 64
-    num_redshift = 64#512#128#64#512#256#256#64#512#128
+    num_redshift = 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)
 
@@ -586,7 +586,7 @@ def train(rank, world_size):
 
     ddp_setup(rank, world_size)
     
-    num_train_image_list = [10]#[200]#[1600,3200,6400,12800,25600]
+    num_train_image_list = [3200]#[200]#[1600,3200,6400,12800,25600]
     for i, num_image in enumerate(num_train_image_list):
         config.num_image = num_image
         # config.world_size = world_size
@@ -677,9 +677,9 @@ if __name__ == "__main__":
     world_size = torch.cuda.device_count()
     print(f" sampling, world_size = {world_size} ".center(100,'-'))
     # num_train_image_list = [1600,3200,6400,12800,25600]
-    num_train_image_list = [2000]
-    num_new_img_per_gpu = 8
-    max_num_img_per_gpu = 2
+    num_train_image_list = [3200]
+    num_new_img_per_gpu = 9
+    max_num_img_per_gpu = 1
 
     params = torch.tensor([4.4, 131.341])
 
@@ -690,7 +690,7 @@ if __name__ == "__main__":
 
     for num_image in num_train_image_list:
         config.num_image = num_image
-        config.resume = f"./outputs/model_state-N{num_image}-epoch3-device0"
+        config.resume = f"./outputs/model_state-N{num_image}-epoch7-device0"
 
         # print("ddpm21cm = DDPM21CM(config)")
         manager = mp.Manager()