Open
refactor(losses): replace standalone VAE loss functions with nn.Module subclasses#71
Conversation
… classes - Add src/embkit/losses/base.py: VAELoss abstract base with step_beta() and _kl_divergence() - Rewrite vae_loss.py: MSELoss, BCELoss, BCEWithLogitsLoss, BCEKLWeightedLoss classes - Add LOSS_REGISTRY + get_loss() factory in vae_loss.py - Keep deprecated free-function wrappers for backward compatibility - Update losses/__init__.py to export all new symbols - Update optimize/__init__.py: fit_vae accepts VAELoss instances; fit_net_vae uses BCELoss/BCEWithLogitsLoss directly - Update rna_vae.py: use BCEKLWeightedLoss + step_beta() instead of manual beta warmup - Update vae_estimator.py: fix broken imports; use BCEWithLogitsLoss - Update commands/model.py: use get_loss() registry instead of if-else dispatch - Update tests: cover new class API; update command assertions to use isinstance checks Co-authored-by: kellrott <113868+kellrott@users.noreply.github.com>
…n fit_net_vae Co-authored-by: kellrott <113868+kellrott@users.noreply.github.com>
Copilot
AI
changed the title
refactor(losses): replace free functions with nn.Module-based VAELoss classes
refactor(losses): replace standalone VAE loss functions with nn.Module subclasses
Aug 12, 2026
Copilot created this pull request from a session on behalf of
kellrott
August 12, 2026 23:39
View session
kellrott
reviewed
Aug 13, 2026
kellrott
left a comment
kellrott
left a comment
Contributor
There was a problem hiding this comment.
These loss functions should probably be mentioned in the documentation.
| def bce(recon_x, x, mu, logvar, beta=1.0) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | ||
| """Calculate the VAE loss. | ||
| Used for classification of class probabilities (e.g., MNIST). | ||
| class MSELoss(VAELoss): |
Contributor
There was a problem hiding this comment.
This may be confused with torch.nn.MSELoss, maybe a name like MSEVAELoss
| Returns: | ||
| tuple: Total loss, reconstruction loss and KL divergence loss respectively as floats. | ||
|
|
||
| class BCELoss(VAELoss): |
Contributor
There was a problem hiding this comment.
Similar to MSELoss, this could be mistaken for torch.nn.BCELoss. Maybe a naming pattern like BCEVAELoss would help to differentiate.
| # Registry + factory helper | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
| LOSS_REGISTRY: dict = { |
Contributor
There was a problem hiding this comment.
These are only VAE losses
| # Legacy free-function API (deprecated) | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
| def _deprecated(fn_name: str) -> None: |
Contributor
There was a problem hiding this comment.
Have you verified that all of the deprecated methods have been removed from the unit tests?
Contributor
|
@copilot respond to PR review |
Co-authored-by: kellrott <113868+kellrott@users.noreply.github.com>
Contributor
Author
Addressed the PR review feedback in 025dd7e: VAE loss class names now use explicit |
- encode command: use result[2] (z sample) instead of result[0] (mu) from the encoder's (mu, logvar, z) output tuple - train_netvae: change --min-group-size CLI default from 0 to 2 to match feature_map_intersect's own default, so groups with no overlapping members are filtered out and NetVAE correctly raises ValueError when there is no overlap between RNA features and pathway genes Co-authored-by: kellrott <113868+kellrott@users.noreply.github.com>
☂️ Python Coverage
Overall Coverage
New Files
Modified Files
|
Co-authored-by: kellrott <113868+kellrott@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
embkit.lossesused stateless free functions with inconsistent signatures, pushingbetawarmup state andkl_weightmanagement into every training loop. Replace them with a propernn.Modulehierarchy.New API
VAELoss— abstract base (nn.Module) withforward(recon, x, mu, logvar) → (total, recon_loss, kl_loss), mutablebetaattribute, andstep_beta(kappa, max_beta)warmup helperMSELoss,BCELoss,BCEWithLogitsLoss,BCEKLWeightedLossLOSS_REGISTRY+get_loss(name, **kwargs)factory — eliminates duplicated if-else dispatchCaller updates
rna_vae.py— removes the manualbetavariable; delegates warmup toloss_fn.step_beta()optimize.fit_vae— setsloss.betaper phase when given aVAELoss; still accepts legacy callablesoptimize.fit_net_vae— allocates a singleBCEWithLogitsLossinstance before the epoch loop instead of constructing one per batchcommands/model.py— replaces duplicated if-else blocks withget_loss(loss)vae_estimator.py— fixes broken imports; usesBCEWithLogitsLossBackward compatibility
All previous free functions (
bce,mse,bce_with_logits,bce_kl_weighted,net_vae_loss) are retained as deprecated thin wrappers emittingDeprecationWarning. CLI interface (--loss mse|bce|bce-logit) is unchanged.