feat: add beta parameter for disentanglement

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

numalogic/models/vae/layer.py

@@ -60,3 +60,37 @@ def __init__(
 
  def forward(self, input_: Tensor) -> Tensor:
  return self.relu(self.bnorm(self.conv(input_)))
+
+
+class ConvTransposeBlock(nn.Module):
+ """Basic transpose convolutional block consisting of:
+ - transpose convolutional layer
+ - batch norm
+ - relu activation.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ padding: int = 1,
+ dilation: int = 1,
+ output_padding: int = 0,
+ ):
+ super().__init__()
+ self.convtranspose = nn.ConvTranspose1d(
+ in_channels,
+ out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ output_padding=output_padding,
+ )
+ self.bnorm = nn.BatchNorm1d(out_channels)
+ self.relu = nn.ReLU(inplace=True)
+
+ def forward(self, input_: Tensor) -> Tensor:
+ return self.relu(self.bnorm(self.convtranspose(input_)))

numalogic/models/vae/variants/conv.py

@@ -33,12 +33,10 @@ def __init__(
  n_features: int,
  latent_dim: int,
  conv_channels: Sequence[int] = (16,),
- num_samples: int = 10,
  ):
  super().__init__()
 
  self.seq_len = seq_len
- self.nsamples = num_samples
 
  conv_layer = CausalConvBlock(
  in_channels=n_features,
@@ -101,7 +99,7 @@ def forward(self, x: Tensor) -> tuple[Tensor, Tensor]:
 
 class Decoder(nn.Module):
  """
- Decoder (non-probabilistic) module for Convolutional Variational Autoencoder.
+ Decoder module for Convolutional Variational Autoencoder.
 
  Args:
  ----
@@ -155,7 +153,6 @@ class Conv1dVAE(BaseVAE):
  n_features: num of features
  conv_channels: number of convolutional channels
  latent_dim: latent dimension
- num_samples: number of samples to draw from the latent distribution
 
  Raises
  ------
@@ -170,21 +167,20 @@ def __init__(
  n_features: int,
  latent_dim: int,
  conv_channels: Sequence[int] = (16,),
- num_samples: int = 10,
+ beta: float = 1.0,
  **kwargs,
  ):
  super().__init__(**kwargs)
  self.seq_len = seq_len
  self.z_dim = latent_dim
  self.n_features = n_features
- self.nsamples = num_samples
+ self.beta = beta
 
  self.encoder = Encoder(
  seq_len=seq_len,
  n_features=n_features,
  conv_channels=conv_channels,
  latent_dim=latent_dim,
- num_samples=num_samples,
  )
  self.decoder = Decoder(
  seq_len=seq_len,
@@ -245,4 +241,4 @@ def training_step(self, batch: Tensor, batch_idx: int) -> Tensor:
  on_epoch=True,
  on_step=False,
  )
- return kld_loss + recon_loss
+ return recon_loss + (self.beta * kld_loss)