Parametric Autoencoder

autoencoder

CLASS	DESCRIPTION
Autoencoder	Autoencoder for unsupervised representation learning and reconstruction.

Classes

Autoencoder(*, model: dict[str, ModelType], loss_kwargs: dict | None = None, optimizer_fn: type[torch.optim.Optimizer] | str | None = None, optimizer_kwargs: dict | None = None, lr_scheduler_fn: Any | str | None = None, lr_scheduler_kwargs: dict | None = None, device: Literal['cuda', 'cpu'] = 'cuda')

Bases: Parametric

Autoencoder for unsupervised representation learning and reconstruction.

This class implements a standard autoencoder architecture with separate encoder and decoder networks. The encoder compresses input data into a lower-dimensional latent representation, while the decoder reconstructs the original input from this representation. The model is trained to minimize reconstruction error.

The autoencoder uses mean squared error (MSE) loss for reconstruction \(L = \frac{1}{N} \sum_{i=1}^{N} \| X_i - \hat{X}_i \|^2_2\), where \(X\) is the input data and \(\hat{X}\) is the reconstructed output.

PARAMETER	DESCRIPTION
model	Dictionary containing 'encoder' and 'decoder' model specifications. Both networks are required for autoencoder functionality: 'encoder': Maps input data to latent representation 'decoder': Maps latent representations back to original input space Example: model={"encoder": encoder_net, "decoder": decoder_net} TYPE: ModelType | dict[str, ModelType]
loss_kwargs	Additional keyword arguments to pass to the MSELoss constructor. "reduce" : Literal["sum", "mean"], optional, by default "mean" Reduction method for loss computation. "sign" : float, optional, by default 1.0 Sign multiplier for loss (1.0 for minimization). TYPE: dict, optional, by default None DEFAULT: None
optimizer_fn	Optimizer specification for training. None: Uses Adam optimizer (default) str: Optimizer name from torch.optim (e.g., "AdamW", "SGD") type[torch.optim.Optimizer]: Optimizer class (supports third-party) TYPE: type[torch.optim.Optimizer] | str | None, optional, by default None DEFAULT: None
optimizer_kwargs	Keyword arguments for optimizer instantiation (e.g., {"lr": 1e-3, "weight_decay": 1e-4}). TYPE: dict | None, optional, by default None DEFAULT: None
lr_scheduler_fn	Learning rate scheduler specification. None: No scheduler (default) str: Scheduler name from torch.optim.lr_scheduler (e.g., "ReduceLROnPlateau") type: Scheduler class TYPE: Any | str | None, optional, by default None DEFAULT: None
lr_scheduler_kwargs	Keyword arguments for scheduler instantiation. TYPE: dict | None, optional, by default None DEFAULT: None
device	Computation device. Automatically falls back to "cpu" if CUDA is unavailable. TYPE: Literal['cuda', 'cpu'] DEFAULT: "cuda"

ATTRIBUTE	DESCRIPTION
decoder	Parsed decoder network on the specified device. Also available as models['decoder'] for consistency with multi-model interface. TYPE: Model
loss_fn	Mean squared error loss function for reconstruction. TYPE: Loss

Examples:

>>> import torch
>>> from spectre.parametric import Autoencoder
>>> from spectre.core import Model
>>>
>>> # Create encoder (input -> latent)
>>> encoder = Model(in_features=784, out_features=64, multipliers=[0.5])
>>> # Create decoder (latent -> output)
>>> decoder = Model(in_features=64, out_features=784, multipliers=[2.0])
>>>
>>> # Initialize with model dictionary
>>> autoencoder = Autoencoder(
..      model={"encoder": encoder, "decoder": decoder}
..  )
>>>
>>> # Training data
>>> X = torch.randn(100, 784)
>>> batch = X[:32]  # Single batch
>>> loss = autoencoder.training_step(batch, 0)

METHOD	DESCRIPTION
loss	Compute autoencoder reconstruction loss between input and decoded output.
forward_step	Forward pass through the encoder network.
score_step	Compute reconstruction loss for evaluation without training status checks.
training_step	Compute autoencoder reconstruction loss for a single training batch.

Source code in spectre/parametric/autoencoder.py

def __init__(
    self,
    *,
    model: dict[str, ModelType],
    loss_kwargs: dict | None = None,
    optimizer_fn: type[torch.optim.Optimizer] | str | None = None,
    optimizer_kwargs: dict | None = None,
    lr_scheduler_fn: Any | str | None = None,
    lr_scheduler_kwargs: dict | None = None,
    device: Literal["cuda", "cpu"] = "cuda",
) -> None:
    if not isinstance(model, dict):
        raise TypeError(
            "model must be a dictionary containing 'encoder' and 'decoder' keys."
        )
    if "encoder" not in model or "decoder" not in model:
        raise ValueError(
            "model dict must contain both 'encoder' and 'decoder' keys."
        )

    super().__init__(
        model=model,
        device=device,
        optimizer_fn=optimizer_fn,
        optimizer_kwargs=optimizer_kwargs,
        lr_scheduler_fn=lr_scheduler_fn,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
    )

    self.decoder = self.model_dict["decoder"]

    if loss_kwargs is None:
        loss_kwargs = {}
    if not isinstance(loss_kwargs, dict):
        raise TypeError(
            "Additional loss parameters `loss_kwargs` must be a dictionary."
        )

    self.loss_fn = MSELoss(**loss_kwargs)

Functions

loss(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute autoencoder reconstruction loss between input and decoded output.

Performs encoding and decoding operations, then computes mean squared error between the original input and the reconstructed output.

PARAMETER	DESCRIPTION
X	Input data tensor to encode and reconstruct. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted loss computation. If None, uniform weighting is applied. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Unused for autoencoders (reconstruction target is always the input). TYPE: torch.Tensor or None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Reconstruction loss value.

Source code in spectre/parametric/autoencoder.py

def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute autoencoder reconstruction loss between input and decoded output.

    Performs encoding and decoding operations, then computes mean squared
    error between the original input and the reconstructed output.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data tensor to encode and reconstruct.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Sample weights for weighted loss computation. If None, uniform
        weighting is applied.

    target : torch.Tensor or None, by default None
        Unused for autoencoders (reconstruction target is always the input).

    Returns
    -------
    torch.Tensor
        Reconstruction loss value.
    """
    z_batch = self.forward_step(X)
    x_batch_pred = self.decoder(z_batch)

    context = {"X": X, "Y": x_batch_pred}
    if weights is not None:
        context["weights"] = weights

    loss = self.loss_fn(context)

    return loss

forward_step(X: torch.Tensor) -> torch.Tensor

Forward pass through the encoder network.

Compresses input data into a lower-dimensional latent representation using the encoder network. This is the first step of the autoencoder pipeline, followed by decoding for reconstruction.

PARAMETER	DESCRIPTION
X	Input data to encode into latent representation. TYPE: torch.Tensor of shape (n_samples, in_features)

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, latent_features)	Encoded latent representation of the input data.

Source code in spectre/parametric/autoencoder.py

def forward_step(self, X: torch.Tensor) -> torch.Tensor:
    """
    Forward pass through the encoder network.

    Compresses input data into a lower-dimensional latent representation
    using the encoder network. This is the first step of the autoencoder
    pipeline, followed by decoding for reconstruction.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data to encode into latent representation.

    Returns
    -------
    torch.Tensor of shape (n_samples, latent_features)
        Encoded latent representation of the input data.
    """
    X = self.model(X)
    return X

score_step(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute reconstruction loss for evaluation without training status checks.

This method provides direct access to the autoencoder's reconstruction loss computation for evaluation purposes, bypassing the is_trained requirement of the inherited score method.

PARAMETER	DESCRIPTION
X	Input data to encode and reconstruct. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted loss computation. If None, uniform weighting is applied. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Unused for autoencoders (reconstruction target is the input itself). TYPE: torch.Tensor or None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Reconstruction loss value.

Source code in spectre/parametric/autoencoder.py

def score_step(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute reconstruction loss for evaluation without training status checks.

    This method provides direct access to the autoencoder's reconstruction
    loss computation for evaluation purposes, bypassing the `is_trained`
    requirement of the inherited `score` method.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data to encode and reconstruct.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Sample weights for weighted loss computation. If None, uniform
        weighting is applied.

    target : torch.Tensor or None, by default None
        Unused for autoencoders (reconstruction target is the input itself).

    Returns
    -------
    torch.Tensor
        Reconstruction loss value.
    """
    return self.loss(X=X, weights=weights, target=target)

training_step(batch: DataBatch, batch_idx: int) -> torch.Tensor

Compute autoencoder reconstruction loss for a single training batch.

Performs forward pass through encoder and decoder networks, then computes reconstruction loss between input and decoded output. Automatically logs the loss for monitoring during training and validation.

PARAMETER	DESCRIPTION
batch	Training batch containing (input_data, sample_weights) where: input_data : torch.Tensor of shape (n_samples, input_features) Input data to encode and reconstruct sample_weights : torch.Tensor of shape (n_samples, ) Sample weights for weighted loss computation TYPE: DataBatch
batch_idx	Batch index. TYPE: int

RETURNS	DESCRIPTION
Tensor	Reconstruction loss value.

Source code in spectre/parametric/autoencoder.py

def training_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    Compute autoencoder reconstruction loss for a single training batch.

    Performs forward pass through encoder and decoder networks, then computes
    reconstruction loss between input and decoded output. Automatically logs
    the loss for monitoring during training and validation.

    Parameters
    ----------
    batch : DataBatch
        Training batch containing (input_data, sample_weights) where:

        - input_data : torch.Tensor of shape (n_samples, input_features)
            Input data to encode and reconstruct
        - sample_weights : torch.Tensor of shape (n_samples, )
            Sample weights for weighted loss computation

    batch_idx : int
        Batch index.

    Returns
    -------
    torch.Tensor
        Reconstruction loss value.
    """
    loss = self.loss(X=batch.data, weights=batch.weights, target=batch.target)

    self.log_training_metrics(loss)

    return loss