Added default parameters to decay rate and steps as well as introduce…

…d use lr decay flag
andreasMazur · Jul 31, 2024 · f0bc9bc · f0bc9bc
1 parent 8d65119
commit f0bc9bc
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/geoconv_examples/mpi_faust/pytorch/model.py b/src/geoconv_examples/mpi_faust/pytorch/model.py
@@ -140,11 +140,12 @@ def train_loop(self,
                    dataset,
                    loss_fn,
                    opt,
-                   decay_rate,
-                   decay_steps,
+                   decay_rate=0.95,
+                   decay_steps=500,
                    verbose=True,
                    epoch=None,
-                   prev_steps=None):
+                   prev_steps=None,
+                   use_lr_decay=False):
         self.train()
         epoch_accuracy = 0.
         epoch_loss = 0.
@@ -158,7 +159,8 @@ def train_loop(self,
             loss.backward()
             opt.step()
 
-            custom_exp_scheduler(opt, prev_steps + step, decay_rate=decay_rate, decay_steps=decay_steps)
+            if use_lr_decay:
+                custom_exp_scheduler(opt, prev_steps + step, decay_rate=decay_rate, decay_steps=decay_steps)
 
             # Statistics
             epoch_accuracy = epoch_accuracy + multiclass_accuracy(pred, gt).detach()