0720-0016

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 = (1, 2, 2, 4, 8)#(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 = (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:

diffusion.py

@@ -231,7 +231,7 @@ class TrainConfig:
     push_to_hub = True 
     hub_model_id = "Xsmos/ml21cm"
     hub_private_repo = False
-    dataset_name = "/storage/home/hcoda1/3/bxia34/scratch/LEN128-DIM64-CUB8-4.4-131.341.h5"
+    dataset_name = "/storage/home/hcoda1/3/bxia34/scratch/LEN128-DIM64-CUB8.h5"
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     world_size = torch.cuda.device_count()
     # repeat = 2
@@ -240,8 +240,8 @@ class TrainConfig:
     dim = 2
     stride = (2,2) if dim == 2 else (2,2,2)
     num_image = 1000#2000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560
-    batch_size = 50#1#2#50#20#2#100 # 10
-    n_epoch = 120#5#4# 10#50#20#20#2#5#25 # 120
+    batch_size = 20#1#2#50#20#2#100 # 10
+    n_epoch = 30#120#5#4# 10#50#20#20#2#5#25 # 120
     HII_DIM = 64
     num_redshift = 64#512#128#64#512#256#256#64#512#128
     channel = 1
@@ -583,7 +583,7 @@ class DDPM21CM:
         return x_last
 # %%
 
-num_train_image_list = [800]
+num_train_image_list = [5000]
 
 def train(rank, world_size):
     config = TrainConfig()
@@ -689,10 +689,10 @@ if __name__ == "__main__":
 
         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:
             print(f" sampling for {params}, world_size = {world_size} ".center(100,'-'))