0711-1241

diffusion.py

@@ -427,7 +427,7 @@ class DDPM21CM:
             pbar_train = tqdm(total=len(self.dataloader), disable=not self.accelerator.is_local_main_process)
             pbar_train.set_description(f"device {torch.cuda.current_device()}, Epoch {ep}")
             for i, (x, c) in enumerate(self.dataloader):
-                print(f"device {torch.cuda.current_device()}, x[:,0,:2,:2,:2] =", x[:,0,:2,:2,:2])
+                print(f"device {torch.cuda.current_device()}, x[:10,0,:2,:2,:2] =", x[:10,0,:2,:2,:2])
                 with self.accelerator.accumulate(self.nn_model):
                     x = x.to(self.config.device)
                     xt, noise, ts = self.ddpm.add_noise(x)

load_h5.py

@@ -102,7 +102,7 @@ class Dataset4h5(Dataset):
             self.images = torch.ones_like(torch.from_numpy(self.images)) * torch.arange(len(self.images))[:,None,None,None,None]
             self.images = self.images.numpy()
             # print("after self.images.shape =", self.images.shape)
-            # print(self.images[:6,0,:2,:2,:2])
+            print(self.images[:6,0,:2,0,0])
             # self.images = self.images.numpy()
             
             # plt.imshow(self.images[0,0,0])