feat: flatten inside vanilla ae

Signed-off-by: Avik Basu <ab93@users.noreply.github.com>

numalogic/config/factory.py

@@ -132,6 +132,7 @@ class ModelFactory(_ObjectFactory):
         SparseLSTMAE,
         TransformerAE,
         SparseTransformerAE,
+        VanillaICAE,
     )
     from numalogic.models.vae.variants import Conv1dVAE
 
@@ -146,6 +147,7 @@ class ModelFactory(_ObjectFactory):
         "TransformerAE": TransformerAE,
         "SparseTransformerAE": SparseTransformerAE,
         "Conv1dVAE": Conv1dVAE,
+        "VanillaICAE": VanillaICAE,
     }
 
 

numalogic/models/autoencoder/variants/icvanilla.py

@@ -5,7 +5,7 @@
 
 from numalogic.models.autoencoder.base import BaseAE
 from numalogic.tools.exceptions import LayerSizeMismatchError
-from numalogic.tools.layer import MultiChannelLinear
+from numalogic.tools.layer import IndependentChannelLinear
 
 
 class _VanillaEncoder(nn.Module):
@@ -52,13 +52,13 @@ def _construct_layers(self, layersizes: Sequence[int]) -> nn.ModuleList:
         start_layersize = self.seq_len
 
         for lsize in layersizes[:-1]:
-            _l = [MultiChannelLinear(start_layersize, lsize, self.n_features)]
+            _l = [IndependentChannelLinear(start_layersize, lsize, self.n_features)]
             if self.bnorm:
                 _l.append(nn.BatchNorm1d(self.n_features))
             layers.extend([*_l, nn.Tanh(), nn.Dropout(p=self.dropout_p)])
             start_layersize = lsize
 
-        _l = [MultiChannelLinear(start_layersize, layersizes[-1], self.n_features)]
+        _l = [IndependentChannelLinear(start_layersize, layersizes[-1], self.n_features)]
         if self.bnorm:
             _l.append(nn.BatchNorm1d(self.n_features))
         layers.extend([*_l, nn.Tanh(), nn.Dropout(p=self.dropout_p)])
@@ -115,23 +115,23 @@ def _construct_layers(self, layersizes: Sequence[int]) -> nn.ModuleList:
         layers = nn.ModuleList()
 
         for idx, _ in enumerate(layersizes[:-1]):
-            _l = [MultiChannelLinear(layersizes[idx], layersizes[idx + 1], self.n_features)]
+            _l = [IndependentChannelLinear(layersizes[idx], layersizes[idx + 1], self.n_features)]
             if self.bnorm:
                 _l.append(nn.BatchNorm1d(self.n_features))
             layers.extend([*_l, nn.Tanh(), nn.Dropout(p=self.dropout_p)])
 
-        layers.append(MultiChannelLinear(layersizes[-1], self.seq_len, self.n_features))
+        layers.append(IndependentChannelLinear(layersizes[-1], self.seq_len, self.n_features))
         return layers
 
 
 class VanillaICAE(BaseAE):
-    r"""Multichannel Vanilla Autoencoder model based on the vanilla encoder and decoder.
+    r"""Vanilla Autoencoder model with Independent isolated Channels based on the vanilla encoder and decoder.
         Each channel is an isolated neural network.
 
     Args:
     ----
         seq_len: sequence length / window length
-        n_features: num of features/channel, each channel is a separate neural network
+        n_channels: num of features/channel, each channel is a separate neural network
         encoder_layersizes: encoder layer size (default = Sequence[int] = (16, 8))
         decoder_layersizes: decoder layer size (default = Sequence[int] = (8, 16))
         dropout_p: the dropout value (default=0.25)
@@ -142,7 +142,7 @@ class VanillaICAE(BaseAE):
     def __init__(
         self,
         seq_len: int,
-        n_features: int = 1,
+        n_channels: int = 1,
         encoder_layersizes: Sequence[int] = (16, 8),
         decoder_layersizes: Sequence[int] = (8, 16),
         dropout_p: float = 0.25,
@@ -152,7 +152,7 @@ def __init__(
         super().__init__(**kwargs)
         self.seq_len = seq_len
         self.dropout_prob = dropout_p
-        self.n_features = n_features
+        self.n_channels = n_channels
 
         if encoder_layersizes[-1] != decoder_layersizes[0]:
             raise LayerSizeMismatchError(
@@ -162,14 +162,14 @@ def __init__(
 
         self.encoder = _VanillaEncoder(
             seq_len=seq_len,
-            n_features=n_features,
+            n_features=n_channels,
             layersizes=encoder_layersizes,
             dropout_p=dropout_p,
             batchnorm=batchnorm,
         )
         self.decoder = _Decoder(
             seq_len=seq_len,
-            n_features=n_features,
+            n_features=n_channels,
             layersizes=decoder_layersizes,
             dropout_p=dropout_p,
             batchnorm=batchnorm,

numalogic/models/autoencoder/variants/vanilla.py

@@ -29,7 +29,7 @@ class _VanillaEncoder(nn.Module):
         n_features: num of features
         layersizes: encoder layer size
         dropout_p: the dropout value
-
+        batchnorm: Flag to enable/diasable batch normalization
     """
 
     def __init__(
@@ -41,13 +41,13 @@ def __init__(
         batchnorm: bool,
     ):
         super().__init__()
-        self.seq_len = seq_len
+        self.input_size = seq_len * n_features
         self.n_features = n_features
         self.dropout_p = dropout_p
         self.bnorm = batchnorm
 
         layers = self._construct_layers(layersizes)
-        self.encoder = nn.Sequential(*layers)
+        self.encoder = nn.Sequential(nn.Flatten(), *layers)
 
     def _construct_layers(self, layersizes: Sequence[int]) -> nn.ModuleList:
         r"""Utility function to generate a simple feedforward network layer.
@@ -61,18 +61,18 @@ def _construct_layers(self, layersizes: Sequence[int]) -> nn.ModuleList:
             A simple feedforward network layer of type nn.ModuleList
         """
         layers = nn.ModuleList()
-        start_layersize = self.seq_len
+        start_layersize = self.input_size
 
         for lsize in layersizes[:-1]:
             _l = [nn.Linear(start_layersize, lsize)]
             if self.bnorm:
-                _l.append(nn.BatchNorm1d(self.n_features))
+                _l.append(nn.BatchNorm1d(lsize))
             layers.extend([*_l, nn.Tanh(), nn.Dropout(p=self.dropout_p)])
             start_layersize = lsize
 
         _l = [nn.Linear(start_layersize, layersizes[-1])]
         if self.bnorm:
-            _l.append(nn.BatchNorm1d(self.n_features))
+            _l.append(nn.BatchNorm1d(layersizes[-1]))
         layers.extend([*_l, nn.Tanh(), nn.Dropout(p=self.dropout_p)])
 
         return layers
@@ -90,7 +90,7 @@ class _Decoder(nn.Module):
         n_features: num of features
         layersizes: decoder layer size
         dropout_p: the dropout value
-
+        batchnorm: flag to enable/disable batch normalization
     """
 
     def __init__(
@@ -102,13 +102,13 @@ def __init__(
         batchnorm: bool,
     ):
         super().__init__()
-        self.seq_len = seq_len
+        self.out_size = seq_len * n_features
         self.n_features = n_features
         self.dropout_p = dropout_p
         self.bnorm = batchnorm
 
         layers = self._construct_layers(layersizes)
-        self.decoder = nn.Sequential(*layers)
+        self.decoder = nn.Sequential(*layers, nn.Unflatten(-1, (n_features, seq_len)))
 
     def forward(self, x: Tensor) -> Tensor:
         return self.decoder(x)
@@ -129,10 +129,10 @@ def _construct_layers(self, layersizes: Sequence[int]) -> nn.ModuleList:
         for idx, _ in enumerate(layersizes[:-1]):
             _l = [nn.Linear(layersizes[idx], layersizes[idx + 1])]
             if self.bnorm:
-                _l.append(nn.BatchNorm1d(self.n_features))
+                _l.append(nn.BatchNorm1d(layersizes[idx + 1]))
             layers.extend([*_l, nn.Tanh(), nn.Dropout(p=self.dropout_p)])
 
-        layers.append(nn.Linear(layersizes[-1], self.seq_len))
+        layers.append(nn.Linear(layersizes[-1], self.out_size))
         return layers
 
 
@@ -315,18 +315,18 @@ def _construct_layers(self, layersizes: Sequence[int]) -> nn.ModuleList:
             A simple feedforward network layer of type nn.ModuleList
         """
         layers = nn.ModuleList()
-        start_layersize = self.seq_len
+        start_layersize = self.input_size
 
         for lsize in layersizes[:-1]:
             _l = [nn.Linear(start_layersize, lsize)]
             if self.bnorm:
-                _l.append(nn.BatchNorm1d(self.n_features))
+                _l.append(nn.BatchNorm1d(lsize))
             layers.extend([*_l, nn.Tanh(), nn.Dropout(p=self.dropout_p)])
             start_layersize = lsize
 
         _l = [nn.Linear(start_layersize, layersizes[-1])]
         if self.bnorm:
-            _l.append(nn.BatchNorm1d(self.n_features))
+            _l.append(nn.BatchNorm1d(layersizes[-1]))
         layers.extend([*_l, nn.ReLU()])
 
         return layers

numalogic/tools/layer.py

@@ -64,15 +64,15 @@ def forward(self, input_: Tensor) -> Tensor:
         return self.relu(self.bnorm(self.conv(input_)))
 
 
-class MultiChannelLinear(nn.Module):
+class IndependentChannelLinear(nn.Module):
     """
     Linear layer that treats each feature as independent isolated channels.
 
     Args:
     ----
         in_features: num of input features
         out_features: num of output features
-        n_channels: num of channels
+        n_channels: num of independent channels
         device: device to run on
         dtype: datatype to use
     """

tests/models/autoencoder/test_trainer.py

@@ -16,6 +16,7 @@
     SparseVanillaAE,
     TransformerAE,
     SparseConv1dAE,
+    VanillaAE,
 )
 from numalogic.tools.data import TimeseriesDataModule, StreamingDataset
 
@@ -93,7 +94,7 @@ def test_trainer_03(self):
         self.assertTupleEqual(self.x_test.shape, y_test.size())
 
     def test_trainer_04(self):
-        model = SparseVanillaAE(seq_len=SEQ_LEN, n_features=self.x_train.shape[1])
+        model = VanillaAE(seq_len=SEQ_LEN, n_features=self.x_train.shape[1], batchnorm=True)
         datamodule = TimeseriesDataModule(SEQ_LEN, self.x_train, batch_size=BATCH_SIZE)
         trainer = TimeseriesTrainer(
             accelerator=ACCELERATOR, max_epochs=EPOCHS, barebones=True, logger=False

tests/models/autoencoder/variants/test_vanilla.py

@@ -75,8 +75,19 @@ def test_multichannel(self):
         test_reconerr = stream_trainer.predict(model, dataloaders=streamloader, unbatch=False)
         self.assertTupleEqual((229, SEQ_LEN, self.X_train.shape[1]), test_reconerr.size())
 
-    def test_vanilla_ic(self):
-        model = VanillaICAE(seq_len=SEQ_LEN, n_features=2)
+    def test_vanilla_ic_01(self):
+        model = VanillaICAE(seq_len=SEQ_LEN, n_channels=2)
+        datamodule = TimeseriesDataModule(SEQ_LEN, self.X_train, batch_size=BATCH_SIZE)
+        trainer = TimeseriesTrainer(fast_dev_run=True, deterministic=True)
+        trainer.fit(model, datamodule=datamodule)
+
+        streamloader = DataLoader(StreamingDataset(self.X_val, SEQ_LEN), batch_size=BATCH_SIZE)
+        stream_trainer = TimeseriesTrainer()
+        test_reconerr = stream_trainer.predict(model, dataloaders=streamloader, unbatch=False)
+        self.assertTupleEqual((229, SEQ_LEN, self.X_train.shape[1]), test_reconerr.size())
+
+    def test_vanilla_ic_02(self):
+        model = VanillaICAE(seq_len=SEQ_LEN, n_channels=2, batchnorm=True, loss_fn="mse")
         datamodule = TimeseriesDataModule(SEQ_LEN, self.X_train, batch_size=BATCH_SIZE)
         trainer = TimeseriesTrainer(fast_dev_run=True, deterministic=True)
         trainer.fit(model, datamodule=datamodule)