diff --git a/docs/cli.md b/docs/cli.md index 1e632a1..ffd769c 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -48,6 +48,7 @@ embkit model train-vae INPUT_PATH [OPTIONS] | `--seed` | `42` | Random seed | | `--bfloat16` | false | Use bfloat16 dtype for reduced memory usage | | `--save-stats` | false | Save training statistics alongside the model | +| `--sampling/--no-sampling` | `true` | Enable reparameterization sampling (VAE). Use `--no-sampling` for a standard autoencoder. | **Examples** diff --git a/src/embkit/commands/model.py b/src/embkit/commands/model.py index 453b0c1..0ee09f0 100644 --- a/src/embkit/commands/model.py +++ b/src/embkit/commands/model.py @@ -38,6 +38,7 @@ @click.option("--zero-mask", default=None, type=float) @click.option("--seed", default=42, type=int) @click.option("--bfloat16", is_flag=True) +@click.option("--sampling/--no-sampling", default=True, show_default=True, help="Enable reparameterization sampling during training (VAE). Use --no-sampling for a standard autoencoder.") def train_vae(input_path: str, group: str, latent: int, @@ -52,7 +53,8 @@ def train_vae(input_path: str, zero_mask: float, save_stats: bool, seed: int, - bfloat16: bool + bfloat16: bool, + sampling: bool ): """ Train VAE model from a TSV file. @@ -110,6 +112,7 @@ def train_vae(input_path: str, latent_dim=latent, encoder_layers=enc_layers_list, decoder_layers=dec_layers_list, + sampling=sampling, device=device, dtype=dtype) loss_func = bce_with_logits @@ -228,6 +231,6 @@ def encode(input_path: str, model_path:str, normalize:str, out:str): m.to(get_device()) result = m.encoder(df_tensor) - martix = result[2].detach().cpu().numpy() - out_df = pd.DataFrame(martix, index=df.index) + matrix = result[0].detach().cpu().numpy() + out_df = pd.DataFrame(matrix, index=df.index) out_df.to_csv(out, sep="\t") diff --git a/src/embkit/models/vae/base_vae.py b/src/embkit/models/vae/base_vae.py index 3486ed6..8e0afe6 100644 --- a/src/embkit/models/vae/base_vae.py +++ b/src/embkit/models/vae/base_vae.py @@ -39,12 +39,14 @@ def __init__(self, features: List[str], encoder: Optional[Encoder] = None, decod def build_encoder(feature_dim: int, latent_dim: int, layers: Optional[LayerList] = None, batch_norm: bool = False, + sampling: bool = True, device=None, dtype=None) -> Encoder: return Encoder( feature_dim=feature_dim, latent_dim=latent_dim, layers=layers, batch_norm=batch_norm, + sampling=sampling, device=device, dtype=dtype ) @@ -73,11 +75,11 @@ def forward(self, x: torch.Tensor): def encode(self, x:torch.Tensor): """ - Run encoder model and return encoded values + Run encoder model and return the latent mean (mu) for stable embeddings. """ with torch.no_grad(): - _, _, z = self.encoder(x) - return z + mu, _, _ = self.encoder(x) + return mu @abstractmethod def to_dict(self): @@ -100,8 +102,8 @@ def __init__(self, encoder) -> None: self.encoder = encoder def forward(self, x): - _, _, z = self.encoder(x) - return z + mu, _, _ = self.encoder(x) + return mu def _import_obj(dotted: str): """Import 'pkg.mod.ClassName' -> object.""" diff --git a/src/embkit/models/vae/encoder.py b/src/embkit/models/vae/encoder.py index b838647..bf64ed1 100644 --- a/src/embkit/models/vae/encoder.py +++ b/src/embkit/models/vae/encoder.py @@ -34,7 +34,7 @@ def __init__(self, batch_norm: bool = False, default_activation: Union[str, None] = "relu", make_latent_heads: bool = True, - sampling : bool = False, + sampling : bool = True, constraint: Optional["NetworkConstraint"] = None, device=None, dtype=None): super().__init__() @@ -113,12 +113,13 @@ def forward(self, x: torch.Tensor): if self._make_latent_heads and (self.z_mean is not None) and (self.z_log_var is not None): mu = self.z_mean(h) logvar = self.z_log_var(h) - if self._sampling: + if self._sampling and self.training: std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) z = mu + eps * std - return mu, logvar, z - return mu, logvar, h + else: + z = mu + return mu, logvar, z return h diff --git a/src/embkit/models/vae/vae.py b/src/embkit/models/vae/vae.py index aa290f6..941eaaf 100644 --- a/src/embkit/models/vae/vae.py +++ b/src/embkit/models/vae/vae.py @@ -25,6 +25,7 @@ def __init__( encoder_layers: Optional[LayerList] = None, decoder_layers: Optional[LayerList] = None, batch_norm: bool = False, + sampling: bool = True, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None ): @@ -35,6 +36,8 @@ def __init__( encoder_layers: list of layer configs for Encoder decoder_layers: list of layer configs for Decoder batch_norm: enable encoder batch normalization blocks + sampling: enable reparameterization sampling during the forward pass (VAE). + Set to False for a deterministic autoencoder (mu is used as z). device: torch device used for module initialization dtype: torch dtype used for module initialization """ @@ -52,6 +55,7 @@ def __init__( self._encoder_layers_cfg = encoder_layers self._decoder_layers_cfg = decoder_layers self._batch_norm = batch_norm + self._sampling = sampling self.latent_dim = latent_dim feature_dim = len(features) @@ -63,6 +67,7 @@ def __init__( latent_dim=latent_dim, layers=encoder_layers, batch_norm=batch_norm, + sampling=sampling, device=device, dtype=dtype ) self.decoder = self.build_decoder( @@ -87,7 +92,8 @@ def to_dict(self): "latent_dim": self.latent_dim, "encoder_layers": self._layers_to_dict(self._encoder_layers_cfg), "decoder_layers": self._layers_to_dict(self._decoder_layers_cfg), - "batch_norm": self._batch_norm + "batch_norm": self._batch_norm, + "sampling": self._sampling, } @classmethod @@ -97,6 +103,7 @@ def from_dict(cls, desc): latent_dim=desc["latent_dim"], encoder_layers=LayerList([Layer.from_dict(li) for li in (desc.get("encoder_layers") or [])]), decoder_layers=LayerList([Layer.from_dict(li) for li in (desc.get("decoder_layers") or [])]), - batch_norm=desc.get("batch_norm", False) + batch_norm=desc.get("batch_norm", False), + sampling=desc.get("sampling", True), ) diff --git a/tests/models/vae_models/test_encoder.py b/tests/models/vae_models/test_encoder.py index 592295e..39e78b6 100644 --- a/tests/models/vae_models/test_encoder.py +++ b/tests/models/vae_models/test_encoder.py @@ -111,6 +111,50 @@ def test_encoder_layer_batch_norm(self): mu, logvar, z = enc(x) assert z.shape == (2, 2) + def test_sampling_enabled_eval_z_same_as_mu(self): + """When sampling=True and eval mode is true, z should not differ from mu, sampling is only on during training.""" + torch.manual_seed(0) + enc = Encoder(feature_dim=8, latent_dim=4, layers=None, sampling=True) + enc.eval() # eval mode disables sampling + x = torch.randn(16, 8) + + mu, logvar, z = enc(x) + + self.assertEqual(mu.shape, z.shape) + # With sampling enabled, z should equal mu (reparameterization is disabled durning eval) + self.assertTrue(torch.allclose(z, mu), + "z should not differ from mu when sampling=True and Eval mode is on") + + def test_sampling_disabled_z_equals_mu(self): + """When sampling=False, z should equal mu (no reparameterization noise).""" + enc = Encoder(feature_dim=8, latent_dim=4, layers=None, sampling=False) + x = torch.randn(16, 8) + + mu, logvar, z = enc(x) + + self.assertEqual(mu.shape, z.shape) + # With sampling disabled, z must be identical to mu + self.assertTrue(torch.allclose(z, mu), + "z should equal mu when sampling=False") + + def test_default_sampling_is_true(self): + """Encoder default should have sampling enabled for proper VAE training.""" + enc = Encoder(feature_dim=6, latent_dim=3) + self.assertTrue(enc._sampling, "Default sampling should be True") + + def test_forward_always_returns_three_tuple_with_latent_heads(self): + """forward() must always return (mu, logvar, z) when make_latent_heads=True.""" + for sampling in (True, False): + enc = Encoder(feature_dim=6, latent_dim=3, sampling=sampling) + x = torch.randn(4, 6) + result = enc(x) + self.assertIsInstance(result, tuple) + self.assertEqual(len(result), 3, + f"Expected 3-tuple with sampling={sampling}") + mu, logvar, z = result + self.assertEqual(mu.shape, z.shape) + self.assertEqual(mu.shape, logvar.shape) + if __name__ == "__main__": unittest.main() \ No newline at end of file