Merge pull request #19 from sparticlesteve/update-ocp-opt-name

Update ocp for opt_name logging fix
mlcommons · Sep 9, 2021 · b796e7a · b796e7a
2 parents 9f27b3d + 32007ec
commit b796e7a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/open_catalyst/configs/mlperf_hpc.yml b/open_catalyst/configs/mlperf_hpc.yml
@@ -52,19 +52,17 @@ optim:
     batch_size: 8
     eval_batch_size: 8
     num_workers: 8
-    # Evaluate every epoch
-    #eval_every: -1000
-    lr_initial: 0.0004
-    lr_gamma: 0.1
 
-    # New LR schedule based on DistributedSampler partitioning
+    optimizer: AdamW
+    lr_initial: 0.0004
+    warmup_steps: 31252 # 4 epochs
+    warmup_factor: 0.2
     lr_milestones:
         - 125008 # 16 epochs
         - 187512 # 24 epochs
         - 250016 # 32 epochs
-    warmup_steps: 31252 # 4 epochs
+    lr_gamma: 0.1
 
-    warmup_factor: 0.2
     max_epochs: 30
     energy_coefficient: 0
     force_coefficient: 50

diff --git a/open_catalyst/ocpmodels/trainers/mlperf_forces_trainer.py b/open_catalyst/ocpmodels/trainers/mlperf_forces_trainer.py
@@ -369,6 +369,7 @@ def train(self):
                            value=self.config["optim"]["batch_size"] * self.config["gpus"])
             mllogger.event(key=mllog.constants.TRAIN_SAMPLES, value=len(self.train_loader.dataset))
             mllogger.event(key=mllog.constants.EVAL_SAMPLES, value=len(self.val_loader.dataset))
+            mllogger.event(key=mllog.constants.OPT_NAME, value=self.config["optim"].get("optimizer", "AdamW"))
             mllogger.event(key=mllog.constants.OPT_BASE_LR, value=self.config["optim"]["lr_initial"])
             mllogger.event(key=mllog.constants.OPT_LR_WARMUP_STEPS, value=self.config["optim"]["warmup_steps"])
             mllogger.event(key=mllog.constants.OPT_LR_WARMUP_FACTOR, value=self.config["optim"]["warmup_factor"])