0719-2124

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, 1, 2, 2, 4, 4)#(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, 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)#
             elif image_size == 32:
                 channel_mult = (1, 2, 2, 4)
             elif image_size == 28:

diffusion.py

@@ -241,7 +241,7 @@ class TrainConfig:
     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 = 20#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 = 64#512#128#64#512#256#256#64#512#128
     channel = 1
@@ -666,13 +666,13 @@ def generate_samples(rank, world_size, config, num_new_img_per_gpu, max_num_img_
 
 if __name__ == "__main__":
     world_size = torch.cuda.device_count()
-    print(f" sampling, world_size = {world_size} ".center(100,'-'))
+    # print(f" sampling, world_size = {world_size} ".center(100,'-'))
     # num_train_image_list = [1600,3200,6400,12800,25600]
     # num_train_image_list = [5000]
     num_new_img_per_gpu = 200
-    max_num_img_per_gpu = 20
+    max_num_img_per_gpu = 40
 
-    params = torch.tensor([4.4, 131.341])
+    # params = torch.tensor([4.4, 131.341])
 
     # print("config = TrainConfig()")
     config = TrainConfig()
@@ -687,7 +687,16 @@ if __name__ == "__main__":
         manager = mp.Manager()
         return_dict = manager.dict()
 
-        mp.spawn(generate_samples, args=(world_size, config, num_new_img_per_gpu, max_num_img_per_gpu, return_dict, params), nprocs=world_size, join=True)
+        params_pairs = [
+            (4.4, 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,'-'))
+            mp.spawn(generate_samples, args=(world_size, config, num_new_img_per_gpu, max_num_img_per_gpu, return_dict, torch.tensor(params)), nprocs=world_size, join=True)
 
         # print("---"*30)
         # print(f"device {torch.cuda.current_device()}, keys = {return_dict.keys()}")