20051216

diffusion.py

@@ -363,7 +363,7 @@ def get_gpu_info(device):
 
 class DDPM21CM:
     def __init__(self, config):
-        config.run_name = datetime.now().strftime("%d%H%M%S") # the unique name of each experiment
+        config.run_name = os.environ.get("SLURM_JOB_ID", datetime.now().strftime("%d%H%M%S")) # the unique name of each experiment
         self.config = config
         self.ddpm = DDPMScheduler(betas=(1e-4, 0.02), num_timesteps=config.num_timesteps, img_shape=config.img_shape, device=config.device, config=config,)#, dtype=config.dtype
 
@@ -381,7 +381,7 @@ class DDPM21CM:
             # print(f"resumed nn_model from {config.resume}")
             self.nn_model.module.load_state_dict(torch.load(config.resume)['unet_state_dict'])
             #self.nn_model.module.to(config.dtype)
-            print(f"{config.run_name} cuda:{torch.cuda.current_device()}/{self.config.global_rank} resumed nn_model from {config.resume} with {sum(x.numel() for x in self.nn_model.parameters())} parameters".center(self.config.str_len,'+'))
+            print(f"{config.run_name} cuda:{torch.cuda.current_device()}/{self.config.global_rank} resumed nn_model from {config.resume} with {sum(x.numel() for x in self.nn_model.parameters())} parameters, {datetime.now().strftime('%d-%H:%M:%S.%f')}".center(self.config.str_len,'+'))
         else:
             print(f"{config.run_name} cuda:{torch.cuda.current_device()}/{self.config.global_rank} initialized nn_model randomly with {sum(x.numel() for x in self.nn_model.parameters())} parameters, {datetime.now().strftime('%d-%H:%M:%S.%f')}".center(self.config.str_len,'+'))
 
@@ -713,6 +713,7 @@ if __name__ == "__main__":
     parser.add_argument("--autocast", type=int, required=False, default=False)
     parser.add_argument("--use_checkpoint", type=int, required=False, default=False)
     parser.add_argument("--dropout", type=float, required=False, default=0)
+    parser.add_argument("--lrate", type=float, required=False, default=1e-4)
 
     args = parser.parse_args()
 
@@ -731,6 +732,7 @@ if __name__ == "__main__":
     config.autocast = bool(args.autocast)
     config.use_checkpoint = bool(args.use_checkpoint)
     config.dropout = args.dropout
+    config.lrate = args.lrate
 
     ############################ training ################################
     if args.train:

perlmutter_diffusion.sbatch

@@ -35,10 +35,11 @@ srun python diffusion.py \
     --num_new_img_per_gpu 4 \
     --max_num_img_per_gpu 2 \
     --gradient_accumulation_steps 1 \
-    --autocast 0 \
+    --autocast 1 \
     --use_checkpoint 1 \
     --dropout 0.1 \
-    --resume ./outputs/model-N1600-device_count4-node4-epoch29-20051216 \
+    --lrate 7e-5 \
+    --resume ./outputs/model-N1600-device_count4-node4-epoch34-20051216 \
     #--train "$SCRATCH/LEN128-DIM64-CUB16-Tvir[4, 6]-zeta[10, 250]-0809-123640.h5" \
 
 date

quantify_results.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a0c13ff1f79531d58b80a67cf6ae9141685d85483eac6006d7fc90f33d55283e
-size 28539850
+oid sha256:8144cb4596999bacefce85cad17af66ea5ebfd5e11d2517361e365fce9895b45
+size 15754478