feat: convolutional vae for multivariate time series (#237)

- Introduce variational autoencoder for multivariate time series data
- Customized trainer for vae
- Causal convolutional modules
- Improved training console logging

Exploration in numaproj/numalogic-benchmarks#3

---------

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

Makefile

@@ -16,7 +16,7 @@ clean:
 	@find . -type f -name "*.py[co]" -exec rm -rf {} +
 
 format: clean
-	poetry run black numalogic/ examples/ tests/ benchmarks/
+	poetry run black numalogic/ examples/ tests/
 
 lint: format
 	poetry run ruff check --fix .

numalogic/models/vae/__init__.py

@@ -0,0 +1,3 @@
+from numalogic.models.vae.trainer import VAETrainer
+
+__all__ = ["VAETrainer"]

numalogic/models/vae/base.py

@@ -0,0 +1,63 @@
+from typing import Callable
+
+import torch.nn.functional as F
+from torch import Tensor, optim
+
+from numalogic.base import TorchModel
+
+
+def _init_criterion(loss_fn: str) -> Callable:
+    if loss_fn == "huber":
+        return F.huber_loss
+    if loss_fn == "l1":
+        return F.l1_loss
+    if loss_fn == "mse":
+        return F.mse_loss
+    raise ValueError(f"Unsupported loss function provided: {loss_fn}")
+
+
+class BaseVAE(TorchModel):
+    """
+    Abstract Base class for all Pytorch based variational autoencoder models.
+
+    Args:
+    ----
+        lr: learning rate (default: 3e-4)
+        weight_decay: weight decay factor weight for regularization (default: 0.0)
+        loss_fn: loss function used to train the model
+                    supported values include: {mse, l1, huber}
+    """
+
+    def __init__(
+        self,
+        lr: float = 3e-4,
+        weight_decay: float = 0.0,
+        loss_fn: str = "mse",
+    ):
+        super().__init__()
+        self._lr = lr
+        self.weight_decay = weight_decay
+        self.criterion = _init_criterion(loss_fn)
+
+    def configure_shape(self, x: Tensor) -> Tensor:
+        """Method to configure the batch shape for each type of model architecture."""
+        return x
+
+    def configure_optimizers(self) -> dict:
+        optimizer = optim.Adam(self.parameters(), lr=self._lr, weight_decay=self.weight_decay)
+        return {"optimizer": optimizer}
+
+    def recon_loss(self, batch: Tensor, recon: Tensor, reduction: str = "sum"):
+        return self.criterion(batch, recon, reduction=reduction)
+
+    def validation_step(self, batch: Tensor, batch_idx: int) -> Tensor:
+        """Validation step for the model."""
+        p, recon = self.forward(batch)
+        loss = self.recon_loss(batch, recon)
+        self.log("val_loss", loss)
+        return loss
+
+    def predict_step(self, batch: Tensor, batch_idx: int, dataloader_idx: int = 0) -> Tensor:
+        """Prediction step for the model."""
+        p, recon = self.forward(batch)
+        return self.recon_loss(batch, recon, reduction="none")

numalogic/models/vae/layer.py

@@ -0,0 +1,62 @@
+from torch import nn, Tensor
+import torch.nn.functional as F
+
+
+class CausalConv1d(nn.Conv1d):
+    """Temporal convolutional layer with causal padding."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+    ):
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=0,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+        self.__padding = (kernel_size - 1) * dilation
+
+    def forward(self, x: Tensor) -> Tensor:
+        return super().forward(F.pad(x, (self.__padding, 0)))
+
+
+class CausalConvBlock(nn.Module):
+    """Basic convolutional block consisting of:
+    - causal 1D convolutional layer
+    - batch norm
+    - relu activation.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        self.conv = CausalConv1d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            dilation=dilation,
+        )
+        self.bnorm = nn.BatchNorm1d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, input_: Tensor) -> Tensor:
+        return self.relu(self.bnorm(self.conv(input_)))

numalogic/models/vae/trainer.py

@@ -0,0 +1,69 @@
+import sys
+import warnings
+from typing import Optional
+
+import torch
+from torch import Tensor
+from pytorch_lightning import Trainer, LightningModule
+
+from numalogic.tools.callbacks import ConsoleLogger
+from numalogic.tools.data import inverse_window
+
+
+class VAETrainer(Trainer):
+    """A PyTorch Lightning Trainer for VAE models.
+
+    Args:
+    ----
+        max_epochs: The maximum number of epochs to train for. (default: 100)
+        logger: Whether to use a console logger to log metrics. (default: True)
+        log_freq: The number of epochs between logging. (default: 5)
+        check_val_every_n_epoch: The number of epochs between validation checks. (default: 5)
+        enable_checkpointing: Whether to enable checkpointing. (default: False)
+        enable_progress_bar: Whether to enable the progress bar. (default: False)
+        enable_model_summary: Whether to enable the model summary. (default: False)
+        **trainer_kw: Additional keyword arguments to pass to the Lightning Trainer.
+    """
+
+    def __init__(
+        self,
+        max_epochs: int = 100,
+        logger: bool = True,
+        log_freq: int = 5,
+        check_val_every_n_epoch: int = 5,
+        enable_checkpointing: bool = False,
+        enable_progress_bar: bool = False,
+        enable_model_summary: bool = False,
+        **trainer_kw
+    ):
+        if not sys.warnoptions:
+            warnings.simplefilter("ignore", category=UserWarning)
+
+        if logger:
+            logger = ConsoleLogger(log_freq=log_freq)
+
+        super().__init__(
+            logger=logger,
+            max_epochs=max_epochs,
+            check_val_every_n_epoch=check_val_every_n_epoch,
+            enable_checkpointing=enable_checkpointing,
+            enable_progress_bar=enable_progress_bar,
+            enable_model_summary=enable_model_summary,
+            **trainer_kw
+        )
+
+    def predict(self, model: Optional[LightningModule] = None, unbatch=True, **kwargs) -> Tensor:
+        r"""Predicts the output of the model.
+
+        Args:
+        ----
+            model: The model to predict with. (default: None)
+            unbatch: Whether to inverse window the output. (default: True)
+            **kwargs: Additional keyword arguments to pass to the Lightning
+                      trainers predict method.
+        """
+        recon_err = super().predict(model, **kwargs)
+        recon_err = torch.vstack(recon_err)
+        if unbatch:
+            return inverse_window(recon_err, method="keep_last")
+        return recon_err

numalogic/models/vae/variants/__init__.py

@@ -0,0 +1,3 @@
+from numalogic.models.vae.variants.conv import Conv1dVAE
+
+__all__ = ["Conv1dVAE"]