06161732

diffusion.py

@@ -272,7 +272,7 @@ class TrainConfig:
     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
+    n_epoch = 200#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
     HII_DIM = 64
     num_redshift = 64#256#512#256#512#256#512#256#512#64#512#64#512#64#256CUDAoom#128#64#512#128#64#512#256#256#64#512#128
     startat = 512-num_redshift
@@ -516,12 +516,12 @@ class DDPM21CM:
         #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")
         if torch.distributed.is_initialized():
-            print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank} torch.distributed.is_initialized")
+            #print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank} torch.distributed.is_initialized")
             torch.distributed.barrier()
         else:
             print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank} torch.distributed.is_initialized False!!!!!!!!!!!!!!!") 
 
-        print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank}; nn_model.device = {self.nn_model.device}")
+        #print(f"cuda:{torch.cuda.current_device()}/{self.config.global_rank}; nn_model.device = {self.nn_model.device}")
         #acc_prep_start = time()
         #self.nn_model, self.optimizer, self.dataloader, self.lr_scheduler = \
         #    self.accelerator.prepare(
@@ -702,18 +702,18 @@ class DDPM21CM:
         # nn_model.train()
         # self.nn_model.to(self.ddpm.device)
 
-        self.accelerator = Accelerator(
-            mixed_precision=self.config.mixed_precision,
-            gradient_accumulation_steps=self.config.gradient_accumulation_steps,
-            log_with="tensorboard",
-            project_dir=os.path.join(self.config.output_dir, "logs"),
-            # distributed_type="MULTI_GPU",
-        )
+        #self.accelerator = Accelerator(
+        #    mixed_precision=self.config.mixed_precision,
+        #    gradient_accumulation_steps=self.config.gradient_accumulation_steps,
+        #    log_with="tensorboard",
+        #    project_dir=os.path.join(self.config.output_dir, "logs"),
+        #    # distributed_type="MULTI_GPU",
+        #)
 
-        self.nn_model, self.optimizer, self.lr_scheduler = \
-            self.accelerator.prepare(
-                self.nn_model, self.optimizer, self.lr_scheduler
-                )
+        #self.nn_model, self.optimizer, self.lr_scheduler = \
+        #    self.accelerator.prepare(
+        #        self.nn_model, self.optimizer, self.lr_scheduler
+        #        )
 
         self.nn_model.eval()
 

load_h5.py

@@ -43,7 +43,7 @@ class Dataset4h5(Dataset):
         dim=2, 
         transform=True, 
         ranges_dict=None, 
-        num_workers=len(os.sched_getaffinity(0))//torch.cuda.device_count(),
+        num_workers=1,#len(os.sched_getaffinity(0))//torch.cuda.device_count(),
         startat=0,
         # shuffle=False,
         ):

perlmutter_diffusion.sbatch

@@ -5,7 +5,7 @@
 #SBATCH -q shared #regular
 #SBATCH -N1 
 #SBATCH --gpus-per-node=1
-#SBATCH -t 0:10:00
+#SBATCH -t 0:59:00
 #SBATCH --ntasks-per-node=1
 #SBATCH -oReport-%j 
 #SBATCH --mail-type=BEGIN,END,FAIL
@@ -39,8 +39,8 @@ srun python diffusion.py \
     --num_image 3200 \
     --batch_size 32 \
     --gradient_accumulation_steps 1 \
-    --num_new_img_per_gpu 50 \
-    --max_num_img_per_gpu 2 \
-    #--resume outputs/model-N2000-device_count1-node8-epoch19-19004529 \
+    --num_new_img_per_gpu 320 \
+    --max_num_img_per_gpu 32 \
+    #--resume outputs/model-N3200-device_count1-node1-epoch29-06121554 \
 
 date