diff --git a/docs/cli.md b/docs/cli.md index d31ae52..e1c27b1 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 | +| `--save-latent` | false | Save latent statistics alongside the model (mu, sigma, losses) | | `--sampling/--no-sampling` | `true` | Enable reparameterization sampling (VAE). Use `--no-sampling` for a standard autoencoder. | For HDF5 input (`--group`), normalization must be `none`. diff --git a/src/embkit/commands/model.py b/src/embkit/commands/model.py index 04b0581..96c90a4 100644 --- a/src/embkit/commands/model.py +++ b/src/embkit/commands/model.py @@ -37,6 +37,7 @@ @click.option("--schedule", "-s", type=str, default=None, help="20:0,20:0.1,40:.3,40:.4") @click.option("--loss", type=click.Choice(["mse", "bce", "bce-logit"]), default="bce-logit") @click.option("--save-stats", is_flag=True) +@click.option("--save-latent", is_flag=True) @click.option("--zero-mask", default=None, type=float) @click.option("--seed", default=42, type=int) @click.option("--bfloat16", is_flag=True) @@ -54,6 +55,7 @@ def train_vae(input_path: str, schedule:str, zero_mask: float, save_stats: bool, + save_latent:bool, seed: int, bfloat16: bool, sampling: bool @@ -143,6 +145,46 @@ def train_vae(input_path: str, stats.to_csv(stats_path, sep="\t") click.echo(f"Stats saved, to {stats_path}") + if save_latent: + vae.eval() # Makes output deterministic. Otherwise, model in training. + losses_df = pd.DataFrame({ + "epoch": range(1, len(vae.history["loss"]) + 1), + "loss": vae.history["loss"], + "recon": vae.history["recon"], + "kl": vae.history["kl"]}) + + losses_df.to_csv(f"{out}.losses_stats.tsv", sep="\t", index=False) + click.echo(f"Training losses saved, to {out}.losses_stats.tsv") + + exportloader = DataLoader(dataset, batch_size=batch_size, shuffle=False) + all_mu = [] + all_logvar = [] + + with torch.no_grad(): + for batch in exportloader: + x_tensor = batch[0] if isinstance(batch, (tuple, list)) else batch + x_tensor = x_tensor.to(device=device, dtype=dtype) + _, mu, logvar, _ = vae(x_tensor) + all_mu.append(mu.cpu()) + all_logvar.append(logvar.cpu()) + + mu = torch.cat(all_mu, dim=0) + logvar = torch.cat(all_logvar, dim=0) + std = torch.sqrt(torch.exp(logvar)) # sigma + + index = df.index if df is not None else None + + mu_np = mu.to(device="cpu", dtype=torch.float32).detach().numpy() + std_np = std.to(device="cpu", dtype=torch.float32).detach().numpy() + mu_df = pd.DataFrame(mu_np, index=index, columns=[f"mu_{i}" for i in range(mu.shape[1])]) + std_df = pd.DataFrame(std_np, index=index, columns=[f"std_{i}" for i in range(std.shape[1])]) + mu_path = f"{out}.latent_mu.tsv" + std_path = f"{out}.latent_std.tsv" + mu_df.to_csv(mu_path, sep="\t") + std_df.to_csv(std_path, sep="\t") + click.echo(f"Latent mu saved, to {mu_path}") + click.echo(f"Latent std (sigma) saved, to {std_path}") + @model.command() @click.argument("input_path", type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str)) diff --git a/tests/commands/test_model.py b/tests/commands/test_model.py index 90f9966..a5ec256 100644 --- a/tests/commands/test_model.py +++ b/tests/commands/test_model.py @@ -128,11 +128,88 @@ def test_train_vae_tsv_branches(self, vae_cls, fit_mock, save_mock): self.assertEqual(result.exit_code, 0, msg=result.output) self.assertIn("No output path provided, using default naming.", result.output) self.assertIn("Stats saved, to vae_latent256_epochs20.model.stats.tsv", result.output) + fit_mock.assert_called_once() self.assertEqual(fit_mock.call_args.kwargs["loss"], model_cmd.mse) self.assertEqual(fit_mock.call_args.kwargs["beta_schedule"], [(0.2, 1), (0.4, 1)]) save_mock.assert_called_once() + def test_train_vae_save_latent(self): + with self.runner.isolated_filesystem(): + with open("rna.tsv", "w", encoding="utf-8") as f: + f.write( + "sample\tG1\tG2\tG3\tG4\n" + "s1\t1\t2\t3\t4\n" + "s2\t4\t3\t2\t1\n" + "s3\t2\t2\t2\t2\n" + ) + # not bflot16... possible training bug? + result = self.runner.invoke( + cli_main, + [ + "model", + "train-vae", + "rna.tsv", + "--epochs", + "2", + "--latent", + "2", + "--encode-layers", + "4", + "--decode-layers", + "4", + "--save-stats", + "--save-latent", + "--out", + "vae.model", + ], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertTrue(Path("vae.model.stats.tsv").exists()) + losses_lines = Path("vae.model.losses_stats.tsv").read_text(encoding="utf-8").splitlines() + self.assertEqual(losses_lines[0].split("\t"), ["epoch", "loss", "recon", "kl"]) + self.assertEqual(len(losses_lines) - 1, 2) + mu_lines = Path("vae.model.latent_mu.tsv").read_text(encoding="utf-8").splitlines() + self.assertEqual(mu_lines[0].split("\t"), ["sample", "mu_0", "mu_1"]) + self.assertEqual([line.split("\t")[0] for line in mu_lines[1:]], ["s1", "s2", "s3"]) + std_lines = Path("vae.model.latent_std.tsv").read_text(encoding="utf-8").splitlines() + std_values = [float(v) for line in std_lines[1:] for v in line.split("\t")[1:]] + self.assertTrue(all(v > 0 for v in std_values)) + + with self.runner.isolated_filesystem(): + writer = H5Writer("matrix.h5", "rna", index=["s1", "s2"], columns=["G1", "G2"]) + writer.set_irow(0, [1.0, 2.0]) + writer.set_irow(1, [3.0, 4.0]) + writer.close() + + result = self.runner.invoke( + cli_main, + [ + "model", + "train-vae", + "matrix.h5", + "--group", + "rna", + "--epochs", + "1", + "--latent", + "2", + "--encode-layers", + "4", + "--decode-layers", + "4", + "--save-latent", + "--out", + "vae_h5.model", + ], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + mu_lines = Path("vae_h5.model.latent_mu.tsv").read_text(encoding="utf-8").splitlines() + self.assertEqual(mu_lines[0].split("\t"), ["", "mu_0", "mu_1"]) + self.assertEqual([line.split("\t")[0] for line in mu_lines[1:]], ["0", "1"]) + @patch.object(model_cmd, "save") @patch.object(model_cmd, "fit_vae") @patch.object(model_cmd, "dataframe_loader", return_value="loader")