04151835

diffusion.py

@@ -267,9 +267,9 @@ class TrainConfig:
     world_size = 1#torch.cuda.device_count()
     # repeat = 2
 
-    # dim = 2
-    dim = 3#2
-    stride = (2,4) if dim == 2 else (2,2,2)
+    dim = 2
+    #dim = 3#2
+    stride = (2,2) if dim == 2 else (2,2,2)
     num_image = 32#0#0#640#320#6400#3000#480#1200#120#3000#300#3000#6000#30#60#6000#1000#2000#20000#15000#7000#25600#3000#10000#1000#10000#5000#2560#800#2560
     batch_size = 1#1#10#50#10#50#20#50#1#2#50#20#2#100 # 10
     n_epoch = 30#50#20#1#50#10#1#50#1#50#5#50#5#50#100#50#100#30#120#5#4# 10#50#20#20#2#5#25 # 120
@@ -504,13 +504,14 @@ class DDPM21CM:
 
 
         # print("!!!!!!!!!!!!!!!!, before prepare, self.dataloader.sampler =", self.dataloader.sampler)
-        print_start = time()
-        print(f"model: {self.nn_model.device}")
+        #model_start = time()
+        #print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank} model: {self.nn_model.device}", f"{time()-model_start:.3f}s")
         #print(f"optimizer: {self.optimizer.state_dict()}")
-        print(f"dataloader: {next(iter(self.dataloader))[0].device}")
-        print(f"lr_scheduler: {self.lr_scheduler.optimizer is self.optimizer}")
-        print_end = time()
-        print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank} print costs {print_end-print_start:.3f}s")
+        #dataloader_start = time()
+        #print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank} dataloader: {next(iter(self.dataloader))[0].device}", f"{time()-dataloader_start:.3f}s")
+        #lr_start = time()
+        #print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank} lr_scheduler: {self.lr_scheduler.optimizer is self.optimizer}", f"{time()-lr_start:.3f}s")
+        #print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank} print costs {print_end-print_start:.3f}s")
 
         acc_prep_start = time()
         self.nn_model, self.optimizer, self.dataloader, self.lr_scheduler = \