bump batch size and dial down logging for plain unet

henryaddison · Aug 5, 2024 · b52c8b7 · b52c8b7
1 parent bbc20c6
commit b52c8b7
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 10 deletions.
diff --git a/src/ml_downscaling_emulator/score_sde_pytorch/configs/deterministic/default_configs.py b/src/ml_downscaling_emulator/score_sde_pytorch/configs/deterministic/default_configs.py
@@ -27,10 +27,11 @@ def get_default_configs():
 
   # training
   config.training = training = ml_collections.ConfigDict()
-  training.batch_size = 16#128
+  training.n_epochs = 20
+  training.batch_size = 16
   training.snapshot_freq = 25
-  training.log_freq = 50
-  training.eval_freq = 1000
+  training.log_freq = 500
+  training.eval_freq = 5000
   ## store additional checkpoints for preemption in cloud computing environments
   training.snapshot_freq_for_preemption = 1000
   ## produce samples at each snapshot.
@@ -41,9 +42,6 @@ def get_default_configs():
   training.random_crop_size = 0
   training.continuous = True
   training.reduce_mean = True
-  training.n_epochs = 20
-  training.snapshot_freq = 5
-  training.eval_freq = 5000
   training.sde = ""
 
   # sampling

diff --git a/...scaling_emulator/score_sde_pytorch/configs/deterministic/ukcp_local_pr_12em_plain_unet.py b/...scaling_emulator/score_sde_pytorch/configs/deterministic/ukcp_local_pr_12em_plain_unet.py
@@ -29,10 +29,8 @@ def get_config():
   # training
   training = config.training
   training.n_epochs = 100
-  training.batch_size = 64
-  training.snapshot_freq = 25
-  training.log_freq = 50
-  training.eval_freq = 1000
+  training.snapshot_freq = 20
+  training.batch_size = 256
 
   # data
   data = config.data