Module Standardize

standardize

CLASS	DESCRIPTION
StandardizeLayer	Neural network layer for data standardization.
WeightedBatchNorm1d	Weighted batch normalization for 1D inputs.

Classes

StandardizeLayer(mean: torch.Tensor, std: torch.Tensor, eps: float = torch.finfo(torch.float32).eps)

Bases: Module

Neural network layer for data standardization.

Applies z-score normalization using precomputed statistics: \((X - \mu) / (\sigma + \epsilon)\), where \(\mu\) is the mean, \(\sigma\) is the standard deviation, and \(\epsilon\) is a small constant for numerical stability. The layer does not learn parameters and is intended for use when data statistics are known in advance.

PARAMETER	DESCRIPTION
mean	Precomputed mean values for normalization. TYPE: Tensor
std	Precomputed standard deviation values for normalization. TYPE: Tensor
eps	Small value added to std for numerical stability. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

METHOD	DESCRIPTION
forward	Apply standardization to input tensor.

Source code in spectre/module/standardize.py

def __init__(
    self,
    mean: torch.Tensor,
    std: torch.Tensor,
    eps: float = torch.finfo(torch.float32).eps,
) -> None:
    super().__init__()

    check_same_shape(mean, std)
    check_same_dtype(mean, std)

    if not isinstance(eps, float):
        raise TypeError(f"`eps` must be a float, got {type(eps)}.")
    check_in_interval(eps, "(0, inf)")

    self.register_buffer("mean", mean)
    self.register_buffer("std", std)
    self.register_buffer("eps", torch.ones_like(std) * eps)

Functions

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

Apply standardization to input tensor.

PARAMETER	DESCRIPTION
X	Input tensor to standardize. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Standardized tensor with same shape as input.

Source code in spectre/module/standardize.py

def forward(self, X: torch.Tensor) -> torch.Tensor:
    """
    Apply standardization to input tensor.

    Parameters
    ----------
    X : torch.Tensor
        Input tensor to standardize.

    Returns
    -------
    torch.Tensor
        Standardized tensor with same shape as input.
    """
    return (X - self.mean) / (self.std + self.eps)

WeightedBatchNorm1d(in_features: int, momentum: float = 0.1, eps: float = torch.finfo(torch.float32).eps, affine: bool = True, track_running_stats: bool = True, dtype: torch.dtype = torch.float32, device: torch.device | str | None = None)

Bases: Module

Weighted batch normalization for 1D inputs.

Extends standard batch normalization to support sample weighting during training. In training mode, computes weighted statistics and updates running statistics. In evaluation mode, uses running statistics for normalization.

PARAMETER	DESCRIPTION
in_features	Number of features in input tensor. TYPE: int
momentum	Momentum for running statistics update. TYPE: float, optional, by default 0.1 DEFAULT: 0.1
eps	Small value for numerical stability. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps
affine	Whether to learn affine transformation parameters. TYPE: bool, optional, by default True DEFAULT: True
track_running_stats	Whether to track running mean and variance. TYPE: bool, optional, by default True DEFAULT: True
dtype	TYPE: dtype DEFAULT: float32
dtype	TYPE: dtype DEFAULT: float32

METHOD	DESCRIPTION
reset_running_stats	Reset running statistics to initial values.
reset_parameters	Reset learnable parameters to initial values.
forward	Apply weighted batch normalization.

Source code in spectre/module/standardize.py

def __init__(
    self,
    in_features: int,
    momentum: float = 0.1,
    eps: float = torch.finfo(torch.float32).eps,
    affine: bool = True,
    track_running_stats: bool = True,
    dtype: torch.dtype = torch.float32,
    device: torch.device | str | None = None,
) -> None:
    super().__init__()

    if not isinstance(in_features, int):
        raise TypeError(
            f"Expected `in_features` to be an integer, "
            f"got {type(in_features).__name__}."
        )
    check_in_interval(in_features, "[1, inf)")
    self.in_features = in_features

    if not isinstance(momentum, float):
        raise TypeError(f"`momentum` must be a float, got {type(momentum)}.")
    check_in_interval(momentum, "[0.0, 1.0]")
    self.momentum = momentum

    if not isinstance(eps, float):
        raise TypeError(f"`eps` must be a float, got {type(eps)}.")
    check_in_interval(eps, "[0.0, inf)")
    self.eps = eps

    if not isinstance(affine, bool):
        raise TypeError(f"`affine` must be a bool, got {type(affine)}.")
    self.affine = affine

    if not isinstance(track_running_stats, bool):
        raise TypeError(
            f"`track_running_stats` must be a bool, got {type(track_running_stats)}."
        )
    self.track_running_stats = track_running_stats

    self.dtype = dtype

    if device is None:
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        self.device = torch.device(device)

    if self.affine:
        self.weight = nn.Parameter(
            torch.ones(in_features, dtype=self.dtype, device=self.device)
        )
        self.bias = nn.Parameter(
            torch.zeros(in_features, dtype=self.dtype, device=self.device)
        )
    else:
        self.register_parameter("weight", None)
        self.register_parameter("bias", None)

    if self.track_running_stats:
        self.register_buffer(
            "running_mean",
            torch.zeros(in_features, dtype=self.dtype, device=self.device),
        )
        self.register_buffer(
            "running_var",
            torch.ones(in_features, dtype=self.dtype, device=self.device),
        )
        self.register_buffer(
            "num_batches_tracked",
            torch.tensor(0, dtype=torch.long, device=self.device),
        )
    else:
        self.register_buffer("running_mean", None)
        self.register_buffer("running_var", None)
        self.register_buffer("num_batches_tracked", None)

    self.reset_running_stats()
    self.reset_parameters()

Functions

reset_running_stats() -> None

Reset running statistics to initial values.

Source code in spectre/module/standardize.py

def reset_running_stats(self) -> None:
    """Reset running statistics to initial values."""
    if self.track_running_stats:
        self.running_mean.zero_()
        self.running_var.fill_(1)
        self.num_batches_tracked.zero_()

reset_parameters() -> None

Reset learnable parameters to initial values.

Source code in spectre/module/standardize.py

def reset_parameters(self) -> None:
    """Reset learnable parameters to initial values."""
    if self.affine:
        nn.init.ones_(self.weight)
        nn.init.zeros_(self.bias)
    self.reset_running_stats()

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

Apply weighted batch normalization.

In training mode (train=True), uses weighted batch statistics and updates running statistics. In evaluation mode (train=False), uses running statistics (weights are ignored).

PARAMETER	DESCRIPTION
X	Input tensor of shape (batch_size, in_features). TYPE: Tensor
weights	Sample weights of shape (batch_size,), if None, uniform weights are used. TYPE: torch.Tensor, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Normalized tensor with same shape as input.

Source code in spectre/module/standardize.py

def forward(
    self, X: torch.Tensor, weights: torch.Tensor | None = None
) -> torch.Tensor:
    """
    Apply weighted batch normalization.

    In training mode (`train=True`), uses weighted batch statistics and updates
    running statistics. In evaluation mode (`train=False`), uses running statistics
    (weights are ignored).

    Parameters
    ----------
    X : torch.Tensor
        Input tensor of shape (batch_size, in_features).

    weights : torch.Tensor, optional, by default None
        Sample weights of shape (batch_size,), if None, uniform weights are used.

    Returns
    -------
    torch.Tensor
        Normalized tensor with same shape as input.
    """
    check_2d(X)

    if X.size(1) != self.in_features:
        raise ValueError(f"Expected {self.in_features} features, got {X.size(1)}.")

    if self.training and self.track_running_stats:
        return self._forward_training(X, weights)
    else:
        return self._forward_eval(X)

Functions