feat: support weight decay in optimizers (#161)

Signed-off-by: Avik Basu <[email protected]>

numalogic/models/autoencoder/base.py

@@ -19,15 +19,30 @@
 
 
 class BaseAE(pl.LightningModule, metaclass=ABCMeta):
- """
+ r"""
  Abstract Base class for all Pytorch based autoencoder models for time-series data.
+
+ Args:
+ loss_fn: loss function used to train the model
+ supported values include: {huber, l1, mae}
+ optim_algo: optimizer algo to be used for training
+ supported values include: {adam, adagrad, rmsprop}
+ lr: learning rate (default: 1e-3)
+ weight_decay: weight decay factor weight for regularization (default: 0.0)
  """
 
- def __init__(self, loss_fn: str = "huber", optim_algo: str = "adam", lr: float = 1e-3):
+ def __init__(
+ self,
+ loss_fn: str = "huber",
+ optim_algo: str = "adam",
+ lr: float = 1e-3,
+ weight_decay: float = 0.0,
+ ):
  super().__init__()
  self.lr = lr
  self.optim_algo = optim_algo
  self.criterion = self.init_criterion(loss_fn)
+ self.weight_decay = weight_decay
 
  self._total_train_loss = 0.0
  self._total_val_loss = 0.0
@@ -58,11 +73,11 @@ def init_criterion(loss_fn: str):
 
  def init_optimizer(self, optim_algo: str):
  if optim_algo == "adam":
- return optim.Adam(self.parameters(), lr=self.lr)
+ return optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
  if optim_algo == "adagrad":
- return optim.Adagrad(self.parameters(), lr=self.lr)
+ return optim.Adagrad(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
  if optim_algo == "rmsprop":
- return optim.RMSprop(self.parameters(), lr=self.lr)
+ return optim.RMSprop(self.parameters(), lr=self.lr, weight_decay=self.weight_decay)
  raise NotImplementedError(f"Unsupported optimizer value provided: {optim_algo}")
 
  def configure_shape(self, batch: Tensor) -> Tensor:

tests/models/autoencoder/variants/test_conv.py

@@ -53,6 +53,8 @@ def test_conv1d_2(self):
  enc_channels=[8, 16, 4],
  enc_kernel_sizes=[3, 3, 3],
  dec_activation="sigmoid",
+ weight_decay=1e-3,
+ optim_algo="rmsprop",
  )
  datamodule = TimeseriesDataModule(SEQ_LEN, self.X_train, batch_size=BATCH_SIZE)
  trainer = AutoencoderTrainer(

tests/models/autoencoder/variants/test_lstm.py

@@ -36,7 +36,7 @@ def setUpClass(cls) -> None:
  cls.X_val = scaler.transform(df[-240:])
 
  def test_lstm_ae(self):
- model = LSTMAE(seq_len=SEQ_LEN, no_features=2, embedding_dim=15)
+ model = LSTMAE(seq_len=SEQ_LEN, no_features=2, embedding_dim=15, weight_decay=1e-3)
  datamodule = TimeseriesDataModule(SEQ_LEN, self.X_train, batch_size=BATCH_SIZE)
  trainer = AutoencoderTrainer(
  accelerator=ACCELERATOR, fast_dev_run=True, enable_progress_bar=True

tests/models/autoencoder/variants/test_transformers.py

@@ -71,6 +71,7 @@ def test_train(self):
  dim_feedforward=64,
  num_encoder_layers=3,
  num_decoder_layers=1,
+ weight_decay=1e-3,
  )
  optimizer = torch.optim.Adam(model.parameters(), lr=LR)
  criterion = nn.HuberLoss(delta=0.5)

tests/models/autoencoder/variants/test_vanilla.py

@@ -35,7 +35,7 @@ def setUpClass(cls) -> None:
  cls.X_val = scaler.transform(df[-240:])
 
  def test_vanilla(self):
- model = VanillaAE(seq_len=SEQ_LEN, n_features=self.X_train.shape[1])
+ model = VanillaAE(seq_len=SEQ_LEN, n_features=self.X_train.shape[1], weight_decay=1e-3)
  datamodule = TimeseriesDataModule(SEQ_LEN, self.X_train, batch_size=BATCH_SIZE)
  trainer = AutoencoderTrainer(fast_dev_run=True, enable_progress_bar=True)
  trainer.fit(model, datamodule=datamodule)