0710-1257

diffusion.py

@@ -240,7 +240,7 @@ class TrainConfig:
     batch_size = 2#2#50#20#2#100 # 10
     n_epoch = 2#10#50#20#20#2#5#25 # 120
     HII_DIM = 28#64
-    num_redshift = 4#128#64#512#256#256#64#512#128
+    num_redshift = 2#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)
 
@@ -359,7 +359,14 @@ class DDPM21CM:
         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)
         # self.shape_loaded = dataset.images.shape
         # print("shape_loaded =", self.shape_loaded)
-        self.dataloader = DataLoader(dataset, batch_size=self.config.batch_size, shuffle=True, num_workers=len(os.sched_getaffinity(0)), pin_memory=True)
+        self.dataloader = DataLoader(
+            dataset, 
+            batch_size=self.config.batch_size, 
+            shuffle=True, 
+            num_workers=1,#len(os.sched_getaffinity(0)), 
+            pin_memory=True,
+            persistent_workers=True,
+            )
         # del dataset
         # self.accelerate(self.config)
         del dataset
@@ -547,7 +554,7 @@ def single_main(rank, world_size):
 if __name__ == "__main__":
     # torch.multiprocessing.set_start_method("spawn")
     # args = (config, nn_model, ddpm, optimizer, dataloader, lr_scheduler)
-    world_size = 1#torch.cuda.device_count()
+    world_size = 2#torch.cuda.device_count()
 
     mp.spawn(single_main, args=(world_size,), nprocs=world_size)
     # notebook_launcher(ddpm21cm.train, num_processes=1, mixed_precision='fp16')