Add LR scheduling patience option to MQF2 (#2291)

awslabs · Sep 12, 2022 · c831753 · c831753
1 parent 33a63c1
commit c831753
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/gluonts/torch/model/deepar/lightning_module.py b/src/gluonts/torch/model/deepar/lightning_module.py
@@ -41,7 +41,7 @@ class DeepARLightningModule(pl.LightningModule):
     weight_decay
         Weight decay regularization parameter, default: ``1e-8``.
     patience
-        Patience parameter for learning rate scheduler.
+        Patience parameter for learning rate scheduler, default: ``10``.
     """
 
     def __init__(

diff --git a/src/gluonts/torch/model/mqf2/lightning_module.py b/src/gluonts/torch/model/mqf2/lightning_module.py
@@ -41,6 +41,8 @@ class MQF2MultiHorizonLightningModule(DeepARLightningModule):
         Learning rate
     weight_decay
         Weight decay during training
+    patience
+        Patience parameter for learning rate scheduler, default: ``10``.
     """
 
     def __init__(
@@ -49,12 +51,14 @@ def __init__(
         loss: DistributionLoss = EnergyScore(),
         lr: float = 1e-3,
         weight_decay: float = 1e-8,
+        patience: int = 10,
     ) -> None:
         super().__init__(
             model=model,
             loss=loss,
             lr=lr,
             weight_decay=weight_decay,
+            patience=patience,
         )
 
     def _compute_loss(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: