0712-1610

diffusion.py

@@ -150,7 +150,7 @@ class DDPMScheduler(nn.Module):
         # for i in range(self.num_timesteps, 0, -1):
         # print(f'sampling!!!')
         pbar_sample = tqdm(total=self.num_timesteps)
-        pbar_sample.set_description("Sampling")
+        pbar_sample.set_description(f"device {torch.cuda.current_device()} sampling")
         for i in reversed(range(0, self.num_timesteps)):
             # print(f'sampling timestep {i:4d}',end='\r')
             t_is = torch.tensor([i]).to(device)
@@ -517,7 +517,7 @@ class DDPM21CM:
             params_backup = params.numpy().copy()
             params = self.rescale(params, self.ranges_dict['params'], to=[0,1])
 
-        print(f"sampling {num_new_img} images with normalized params = {params}")
+        print(f"device {torch.cuda.current_device()} sampling {num_new_img} images with normalized params = {params}")
         params = params.repeat(num_new_img,1)
         assert params.dim() == 2, "params must be a 2D torch.tensor"
         # print("params =", params)
@@ -640,7 +640,7 @@ if __name__ == "__main__":
     world_size = torch.cuda.device_count()
     # num_image_list = [1600,3200,6400,12800,25600]
     num_image_list = [100]
-    num_new_img = 12
+    num_new_img = 6
     max_num_img_per_gpu = 2
 
     # print("config = TrainConfig()")