0721-1845

diffusion.py

@@ -241,9 +241,9 @@ 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 = 100#30#120#5#4# 10#50#20#20#2#5#25 # 120
+    n_epoch = 50#100#30#120#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
+    num_redshift = 64#256CUDAoom#128#64#512#128#64#512#256#256#64#512#128
     channel = 1
     img_shape = (channel, HII_DIM, num_redshift) if dim == 2 else (channel, HII_DIM, HII_DIM, num_redshift)
 
@@ -366,18 +366,28 @@ class DDPM21CM:
         self.ranges_dict = config.ranges_dict
 
     def load(self):
-        dataset = Dataset4h5(self.config.dataset_name, num_image=self.config.num_image, HII_DIM=self.config.HII_DIM, num_redshift=self.config.num_redshift, drop_prob=self.config.drop_prob, dim=self.config.dim, ranges_dict=self.ranges_dict)
+        # rank = torch.cuda.current_device()
+        dataset = Dataset4h5(
+            self.config.dataset_name, 
+            num_image=self.config.num_image,
+            idx = 'random',
+            HII_DIM=self.config.HII_DIM, 
+            num_redshift=self.config.num_redshift,
+            drop_prob=self.config.drop_prob, 
+            dim=self.config.dim,
+            ranges_dict=self.ranges_dict
+            )
         # self.shape_loaded = dataset.images.shape
         # print("shape_loaded =", self.shape_loaded)
         # print(f"load, current_device() = {torch.cuda.current_device()}")
         self.dataloader = DataLoader(
             dataset=dataset, 
             batch_size=self.config.batch_size, 
-            shuffle=False, 
-            num_workers=1,#len(os.sched_getaffinity(0)), 
+            shuffle=True,#False, 
+            num_workers=len(os.sched_getaffinity(0)), 
             pin_memory=True,
             persistent_workers=True,
-            sampler=DistributedSampler(dataset),
+            # sampler=DistributedSampler(dataset),
             )
 
         del dataset
@@ -414,14 +424,14 @@ class DDPM21CM:
 
 
         # print("!!!!!!!!!!!!!!!!, before prepare, self.dataloader.sampler =", self.dataloader.sampler)
-        # self.nn_model, self.optimizer, self.dataloader, self.lr_scheduler = \
-        #     self.accelerator.prepare(
-        #     self.nn_model, self.optimizer, self.dataloader, self.lr_scheduler
-        #     )
-        self.nn_model, self.optimizer, self.lr_scheduler = \
+        self.nn_model, self.optimizer, self.dataloader, self.lr_scheduler = \
             self.accelerator.prepare(
-            self.nn_model, self.optimizer, self.lr_scheduler
+            self.nn_model, self.optimizer, self.dataloader, self.lr_scheduler
             )
+        # self.nn_model, self.optimizer, self.lr_scheduler = \
+        #     self.accelerator.prepare(
+        #     self.nn_model, self.optimizer, self.lr_scheduler
+        #     )
 
         # print("!!!!!!!!!!!!!!!!, after prepare, self.dataloader.sampler =", self.dataloader.sampler)
         # print("!!!!!!!!!!!!!!!!, after prepare, self.dataloader.batch_sampler =", self.dataloader.batch_sampler)
@@ -430,7 +440,7 @@ class DDPM21CM:
         global_step = 0
         for ep in range(self.config.n_epoch):
             self.ddpm.train()
-            self.dataloader.sampler.set_epoch(ep)
+            # self.dataloader.sampler.set_epoch(ep)
 
             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}")
@@ -527,7 +537,7 @@ class DDPM21CM:
         # n_sample = params.shape[0]
         # file = self.config.resume
 
-        print(f"device {torch.cuda.current_device()}, sample, params = {params}")
+        # print(f"device {torch.cuda.current_device()}, sample, params = {params}")
         if params is None:
             params = torch.tensor([4.4, 131.341])
             # params_backup = params.numpy().copy()
@@ -583,7 +593,7 @@ class DDPM21CM:
         return x_last
 # %%
 
-num_train_image_list = [8000]
+num_train_image_list = [200]#[8000]
 
 def train(rank, world_size):
     config = TrainConfig()

load_h5.py

@@ -26,14 +26,28 @@ import datetime
 # from huggingface_hub import create_repo, upload_folder
 
 class Dataset4h5(Dataset):
-    def __init__(self, dir_name, num_image=10, field='brightness_temp', shuffle=False, idx=None, num_redshift=512, HII_DIM=64, rescale=True, drop_prob = 0, dim=2, transform=True, ranges_dict=None):
+    def __init__(
+        self,
+        dir_name, 
+        num_image=10, 
+        field='brightness_temp', 
+        idx=None, 
+        num_redshift=512, 
+        HII_DIM=64, 
+        rescale=True, 
+        drop_prob = 0, 
+        dim=2, 
+        transform=True, 
+        ranges_dict=None, 
+        # shuffle=False,
+        ):
         super().__init__()
         
         self.dir_name = dir_name
         self.num_image = num_image
-        self.field = field
-        self.shuffle = shuffle
         self.idx = idx
+        self.field = field
+        # self.shuffle = shuffle
         self.num_redshift = num_redshift
         self.HII_DIM = HII_DIM
         self.drop_prob = drop_prob
@@ -81,14 +95,24 @@ class Dataset4h5(Dataset):
             self.params_keys = list(f['params']['keys'])
             print(f"params keys = {self.params_keys}")
 
-            if self.idx is None:
-                if self.shuffle:
-                    self.idx = np.sort(random.sample(range(max_num_image), self.num_image))
-                    print(f"loading {self.num_image} images randomly")
-                    # print(self.idx)
-                else:
-                    self.idx = range(self.num_image)
-                    print(f"loading {len(self.idx)} images with idx = {self.idx}")
+            # if self.idx is None:
+            #     if self.shuffle:
+            #         self.idx = np.sort(random.sample(range(max_num_image), self.num_image))
+            #         print(f"loading {self.num_image} images randomly")
+            #         # print(self.idx)
+            #     else:
+            #         self.idx = range(self.num_image)
+            #         print(f"loading {len(self.idx)} images with idx = {self.idx}")
+            if self.idx == "random":
+                self.idx = np.sort(random.sample(range(max_num_image), self.num_image))
+                print(f"loading {self.num_image} images randomly with idx = {self.idx}")
+                # print(self.idx)
+            elif self.idx == "range":
+                rank = torch.cuda.current_device()
+                self.idx = range(
+                    rank*self.num_image, (rank+1)*self.num_image
+                    )
+                print(f"loading {len(self.idx)} images with idx = {self.idx}")
             else:
                 print(f"loading {len(self.idx)} images with idx = {self.idx}")