0718-1052

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, 4, 4, 4)#(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, 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

@@ -239,7 +239,7 @@ class TrainConfig:
     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 = 1#2#50#20#2#100 # 10
-    n_epoch = 8#4# 10#50#20#20#2#5#25 # 120
+    n_epoch = 2#4# 10#50#20#20#2#5#25 # 120
     HII_DIM = 64
     num_redshift = 512#128#64#512#256#256#64#512#128
     channel = 1
@@ -586,7 +586,7 @@ def train(rank, world_size):
 
     ddp_setup(rank, world_size)
     
-    num_train_image_list = [200]#[3200]#[200]#[1600,3200,6400,12800,25600]
+    num_train_image_list = [3200]#[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