Core Model

model

CLASS	DESCRIPTION
Model	Base model class for neural networks with post-training transformations.

Classes

Model(in_features: int | None = None, out_features: int | None = None, multipliers: list[float] | None = None, layers: OrderedDict | torch.nn.Sequential | None = None, activation_fn: Callable = torch.nn.ELU, output_fn: torch.nn.Module | None = None, dropout_prob: float = 0.0, batch_norm: bool = False, orthogonal_weights: bool = False)

Bases: Module, ABC

Base model class for neural networks with post-training transformations.

Creates customizable neural networks with optional batch normalization, dropout, and orthogonal weight parametrization. Supports both layer-based and parameter-based construction with post-training transformation capabilities. Supports JIT compilation.

Unlike Estimator for non-parametric methods, which follows the scikit-learn pattern of fit, predict, and score, Model is a PyTorch torch.nn.Module subclass and follows the PyTorch pattern of defining a forward method. This makes Model suitable for parametric methods such as neural networks.

Model can be used to define simple feedforward networks with its parameters such as in_features, out_features, and multipliers. For more complex architectures, pre-constructed layers can be passed via the layers parameter. Additionally, Model supports adding post-training transformations through the add_transform() method, which can be used to apply operations like normalization or dimensionality reduction after the main network.

PARAMETER	DESCRIPTION
in_features	Number of input features. Required when layers is None. TYPE: int | None, optional, by default None DEFAULT: None
out_features	Number of output features. Required when layers is None. TYPE: int | None, optional, by default None DEFAULT: None
multipliers	Layer width multipliers for automatic construction. TYPE: list[float] | None, optional, by default None DEFAULT: None
layers	Pre-constructed layers to use instead of automatic construction. TYPE: OrderedDict | torch.nn.Sequential | None, optional, by default None DEFAULT: None
activation_fn	Activation function class for hidden layers. TYPE: Callable, by default torch.nn.ELU DEFAULT: ELU
output_fn	Final output activation function, if None no output activation is applied. TYPE: torch.nn.Module | None, optional, by default None DEFAULT: None
dropout_prob	Dropout probability in range [0, 1). Zero disables dropout. TYPE: float, by default 0.0 DEFAULT: 0.0
batch_norm	Whether to use batch normalization after linear layers. TYPE: bool, by default False DEFAULT: False
orthogonal_weights	Whether to use orthogonal weight parametrization. TYPE: bool, by default False DEFAULT: False

ATTRIBUTE	DESCRIPTION
in_features	Number of input features. TYPE: int
out_features	Number of output features. TYPE: int

Examples:

Creating a simple feedforward network with specified layer widths:

>>> import torch
>>> from spectre.core import Model
>>> model = Model(
...     in_features=10,
...     out_features=2,
...     multipliers=[0.8, 0.5],
...     activation_fn=torch.nn.ReLU,
... )
>>> print(model)
Model(
  (sequential): Sequential(
    (layer_0): Linear(in_features=10, out_features=8, bias=True)
    (act_0): ReLU()
    (layer_1): Linear(in_features=8, out_features=5, bias=True)
    (act_1): ReLU()
    (layer_2): Linear(in_features=5, out_features=2, bias=True)
  )
>>> X = torch.randn(5, 10)  # Example input
>>> output = model(X)
>>> print(output.shape)
torch.Size([5, 2])

Creating a model from pre-constructed layers:

>>> layers = torch.nn.Sequential(
...     torch.nn.Linear(10, 20),
...     torch.nn.ReLU(),
...     torch.nn.Linear(20, 5),
...     torch.nn.ReLU(),
...     torch.nn.Linear(5, 2),
... )
>>> model = Model(layers=layers)
>>> print(model)
Model(
  (sequential): Sequential(
    (0): Linear(in_features=10, out_features=20, bias=True)
    (1): ReLU()
    (2): Linear(in_features=20, out_features=5, bias=True)
    (3): ReLU()
    (4): Linear(in_features=5, out_features=2, bias=True)
  )
)
>>> X = torch.randn(5, 10)  # Example input
>>> output = model(X)
>>> print(output.shape)
torch.Size([5, 2])

Adding a post-training transformation:

>>> from spectre.transform import L2Normalizer
>>> model.add_transform(L2Normalizer(), name="normalize")
>>> print(model.get_transforms())
[{'name': 'normalize', 'type': 'Transformer'}]
>>> output = model(X)  # Forward pass with transformation applied
>>> print(output.shape)
torch.Size([5, 2])

Notes

• When using layers, in_features and out_features are inferred from the first and last layers respectively.
• Post-training transformations are applied in the order they were added.
• Transformations can be Transformer instances or any torch.nn.Module.
• Supports JIT compilation via torch.jit.script or torch.jit.trace.

METHOD	DESCRIPTION
forward	Forward pass with optional post-training transformations.
add_transform	Add post-training transformation to the model.
clear_transforms	Clear all post-training transformations.
get_transforms	Get information about registered transformations.
freeze_layers	Freeze layers to prevent gradient updates.
unfreeze_layers	Unfreeze layers to enable gradient updates.
get_frozen_layers	Get names of frozen layers.
layer_shapes	Get input/output shapes for each layer.
count_parameters	Count total parameters in model.
summary	Get model summary string.
compress	Create a compressed version of the model by reducing layer sizes.

Source code in spectre/core/model.py

def __init__(
    self,
    in_features: int | None = None,
    out_features: int | None = None,
    multipliers: list[float] | None = None,
    layers: OrderedDict | torch.nn.Sequential | None = None,
    activation_fn: Callable = torch.nn.ELU,
    output_fn: torch.nn.Module | None = None,
    dropout_prob: float = 0.0,
    batch_norm: bool = False,
    orthogonal_weights: bool = False,
) -> None:
    super(Model, self).__init__()

    self._post_transforms = torch.nn.ModuleList()
    self._transform_names: list[str] = []
    self._transform_types: list[str] = []

    self._validate_inputs(
        in_features, out_features, multipliers, layers, dropout_prob
    )

    if layers is not None:
        self._init_from_layers(layers)
    else:
        self._init_from_params(
            in_features,
            out_features,
            multipliers,
            activation_fn,
            output_fn,
            dropout_prob,
            batch_norm,
            orthogonal_weights,
        )

    self._is_trained = False

Attributes

in_features: int property writable

Number of input features.

out_features: int property writable

Number of output features.

is_trained: bool property writable

Whether the model has been trained.

Functions

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

Forward pass with optional post-training transformations.

PARAMETER	DESCRIPTION
X	Input tensor of shape (n_samples, in_features). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Output tensor of shape (n_samples, out_features).

Source code in spectre/core/model.py

def forward(self, X: torch.Tensor) -> torch.Tensor:
    """
    Forward pass with optional post-training transformations.

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

    Returns
    -------
    torch.Tensor
        Output tensor of shape (n_samples, out_features).
    """
    X = self.sequential(X)

    for transform in self._post_transforms:
        X = transform(X)

    return X

add_transform(transform: Transformer | torch.nn.Module, name: str | None = None) -> None

Add post-training transformation to the model.

Transformations are applied after the main network in the order they were added. Supports both Transformer instances (with transform() and inverse_transform() methods) and standard torch.nn.Module instances.

PARAMETER	DESCRIPTION
transform	Transformation to add. Can be a Transformer instance (e.g., ProjectionTransformer, L2Normalizer) or any torch.nn.Module. TYPE: Transformer | Module
name	Transform name. Auto-generated if None. TYPE: str | None, optional, by default None DEFAULT: None

Examples:

>>> from spectre.transform import ProjectionTransformer, L2Normalizer
>>> model = Model(in_features=10, out_features=5)
>>> # Add projection transformer
>>> proj = ProjectionTransformer(torch.randn(3, 5))
>>> model.add_transform(proj, "projection")
>>> # Add normalization
>>> model.add_transform(L2Normalizer(), "normalize")

Source code in spectre/core/model.py

def add_transform(
    self,
    transform: Transformer | torch.nn.Module,
    name: str | None = None,
) -> None:
    """
    Add post-training transformation to the model.

    Transformations are applied after the main network in the order they
    were added. Supports both [Transformer][spectre.core.Transformer] instances
    (with `transform()` and `inverse_transform()` methods) and standard
    `torch.nn.Module` instances.

    Parameters
    ----------
    transform : Transformer | torch.nn.Module
        Transformation to add. Can be a `Transformer` instance (e.g.,
        `ProjectionTransformer`, `L2Normalizer`) or any `torch.nn.Module`.

    name : str | None, optional, by default None
        Transform name. Auto-generated if None.

    Examples
    --------
    >>> from spectre.transform import ProjectionTransformer, L2Normalizer
    >>> model = Model(in_features=10, out_features=5)
    >>> # Add projection transformer
    >>> proj = ProjectionTransformer(torch.randn(3, 5))
    >>> model.add_transform(proj, "projection")
    >>> # Add normalization
    >>> model.add_transform(L2Normalizer(), "normalize")
    """
    if not isinstance(transform, torch.nn.Module):
        raise TypeError(
            f"transform must be a Transformer or torch.nn.Module, "
            f"got {type(transform).__name__}."
        )

    if name is None:
        name = f"transform_{len(self._post_transforms)}"

    self._post_transforms.append(transform)

    if isinstance(transform, Transformer):
        transform_type = "Transformer"
    else:
        transform_type = "Module"

    self._transform_names.append(name)
    self._transform_types.append(transform_type)

clear_transforms() -> None

Clear all post-training transformations.

Source code in spectre/core/model.py

def clear_transforms(self) -> None:
    """Clear all post-training transformations."""
    self._post_transforms = torch.nn.ModuleList()
    self._transform_names.clear()
    self._transform_types.clear()

get_transforms() -> list[dict[str, str]]

Get information about registered transformations.

RETURNS	DESCRIPTION
list[dict[str, str]]	List of dicts with "name" and "type" keys for each transform.

Source code in spectre/core/model.py

def get_transforms(self) -> list[dict[str, str]]:
    """
    Get information about registered transformations.

    Returns
    -------
    list[dict[str, str]]
        List of dicts with "name" and "type" keys for each transform.
    """
    return [
        {"name": name, "type": type_}
        for name, type_ in zip(self._transform_names, self._transform_types)
    ]

freeze_layers(layer_names: list[str] | None = None, freeze_all: bool = False) -> None

Freeze layers to prevent gradient updates.

PARAMETER	DESCRIPTION
layer_names	Names of layers to freeze. If None, uses freeze_all. TYPE: list[str] | None DEFAULT: None
freeze_all	If True, freeze all layers. TYPE: bool, by default False DEFAULT: False

Examples:

>>> model = Model(in_features=100, out_features=10, multipliers=[0.5])
>>> model.freeze_layers(["layer_0"])  # Freeze first layer
>>> model.freeze_layers(freeze_all=True)  # Freeze all

Source code in spectre/core/model.py

def freeze_layers(
    self,
    layer_names: list[str] | None = None,
    freeze_all: bool = False,
) -> None:
    """
    Freeze layers to prevent gradient updates.

    Parameters
    ----------
    layer_names : list[str] | None, optional
        Names of layers to freeze. If None, uses freeze_all.

    freeze_all : bool, by default False
        If True, freeze all layers.

    Examples
    --------
    >>> model = Model(in_features=100, out_features=10, multipliers=[0.5])
    >>> model.freeze_layers(["layer_0"])  # Freeze first layer
    >>> model.freeze_layers(freeze_all=True)  # Freeze all
    """
    if freeze_all:
        for param in self.parameters():
            param.requires_grad = False
    elif layer_names is not None:
        for name in layer_names:
            if hasattr(self.sequential, name):
                layer = getattr(self.sequential, name)
                for param in layer.parameters():
                    param.requires_grad = False
            else:
                raise ValueError(f"Layer '{name}' not found in model.")

unfreeze_layers(layer_names: list[str] | None = None, unfreeze_all: bool = False) -> None

Unfreeze layers to enable gradient updates.

PARAMETER	DESCRIPTION
layer_names	Names of layers to unfreeze. If None, uses unfreeze_all. TYPE: list[str] | None DEFAULT: None
unfreeze_all	If True, unfreeze all layers. TYPE: bool, by default False DEFAULT: False

Source code in spectre/core/model.py

def unfreeze_layers(
    self,
    layer_names: list[str] | None = None,
    unfreeze_all: bool = False,
) -> None:
    """
    Unfreeze layers to enable gradient updates.

    Parameters
    ----------
    layer_names : list[str] | None, optional
        Names of layers to unfreeze. If None, uses unfreeze_all.

    unfreeze_all : bool, by default False
        If True, unfreeze all layers.
    """
    if unfreeze_all:
        for param in self.parameters():
            param.requires_grad = True
    elif layer_names is not None:
        for name in layer_names:
            if hasattr(self.sequential, name):
                layer = getattr(self.sequential, name)
                for param in layer.parameters():
                    param.requires_grad = True
            else:
                raise ValueError(f"Layer '{name}' not found in model.")

get_frozen_layers() -> list[str]

Get names of frozen layers.

RETURNS	DESCRIPTION
list[str]	Names of layers with all parameters frozen.

Source code in spectre/core/model.py

def get_frozen_layers(self) -> list[str]:
    """
    Get names of frozen layers.

    Returns
    -------
    list[str]
        Names of layers with all parameters frozen.
    """
    frozen = []
    for name, module in self.sequential.named_children():
        params = list(module.parameters())
        if params and all(not p.requires_grad for p in params):
            frozen.append(name)
    return frozen

layer_shapes() -> dict[str, tuple]

Get input/output shapes for each layer.

RETURNS	DESCRIPTION
dict[str, tuple]	Mapping from layer name to (in_shape, out_shape).

Examples:

>>> model = Model(in_features=100, out_features=10, multipliers=[0.5])
>>> model.layer_shapes()
{'layer_0': (100, 50), 'layer_1': (50, 10)}

Source code in spectre/core/model.py

def layer_shapes(self) -> dict[str, tuple]:
    """
    Get input/output shapes for each layer.

    Returns
    -------
    dict[str, tuple]
        Mapping from layer name to (in_shape, out_shape).

    Examples
    --------
    >>> model = Model(in_features=100, out_features=10, multipliers=[0.5])
    >>> model.layer_shapes()
    {'layer_0': (100, 50), 'layer_1': (50, 10)}
    """
    shapes = {}
    for name, module in self.sequential.named_children():
        if hasattr(module, "in_features") and hasattr(module, "out_features"):
            shapes[name] = (module.in_features, module.out_features)
    return shapes

count_parameters(trainable_only: bool = False) -> int

Count total parameters in model.

PARAMETER	DESCRIPTION
trainable_only	If True, only count trainable parameters. TYPE: bool, by default False DEFAULT: False

RETURNS	DESCRIPTION
int	Number of parameters.

Source code in spectre/core/model.py

def count_parameters(self, trainable_only: bool = False) -> int:
    """
    Count total parameters in model.

    Parameters
    ----------
    trainable_only : bool, by default False
        If True, only count trainable parameters.

    Returns
    -------
    int
        Number of parameters.
    """
    if trainable_only:
        return sum(p.numel() for p in self.parameters() if p.requires_grad)
    return sum(p.numel() for p in self.parameters())

summary() -> str

Get model summary string.

RETURNS	DESCRIPTION
str	Formatted summary with layer info and parameter counts.

Source code in spectre/core/model.py

def summary(self) -> str:
    """
    Get model summary string.

    Returns
    -------
    str
        Formatted summary with layer info and parameter counts.
    """
    total_params = self.count_parameters()
    trainable_params = self.count_parameters(trainable_only=True)

    lines = [
        f"Model Summary: {type(self).__name__}",
        f"  Input features:  {self.in_features}",
        f"  Output features: {self.out_features}",
        f"  Total layers:    {len(self)}",
        f"  Total params:    {total_params:,}",
        f"  Trainable:       {trainable_params:,}",
        "  Layer-wise shapes:",
    ]

    for name, shape in self.layer_shapes().items():
        lines.append(f"  {name:20s} {shape[0]:6d} -> {shape[1]:6d}")

    lines.append("=" * 60)
    return "\n".join(lines)

compress(ratio: float, preserve_io: bool = True) -> Model

Create a compressed version of the model by reducing layer sizes.

Creates a new model with reduced hidden layer dimensions while optionally preserving input and output dimensions. Only works with sequential models containing at least two linear layers.

PARAMETER	DESCRIPTION
ratio	Compression ratio between 0 and 1. For example, ratio=0.5 creates a model with half the hidden layer sizes. TYPE: float
preserve_io	Whether to preserve input and output dimensions. If True, only compresses hidden layers. If False, compresses all layers except the final output layer. TYPE: bool, by default True DEFAULT: True

RETURNS	DESCRIPTION
Model	New compressed model with reduced layer sizes.

RAISES	DESCRIPTION
TypeError	If ratio is not a float.
ValueError	If ratio is not in the range (0, 1).

Examples:

>>> model = Model(in_features=100, out_features=10, multipliers=[2.0, 1.5])
>>> print(model.layer_shapes())
{'layer_0': (100, 200), 'layer_1': (200, 150), 'layer_2': (150, 10)}
>>> compressed = model.compress(ratio=0.5, preserve_io=True)
>>> print(compressed.layer_shapes())
{'layer_0': (100, 100), 'layer_1': (100, 75), 'layer_2': (75, 10)}

Compress all layers except output:

>>> compressed = model.compress(ratio=0.5, preserve_io=False)
>>> print(compressed.layer_shapes())
{'layer_0': (50, 100), 'layer_1': (100, 75), 'layer_2': (75, 10)}

Notes

• The original model is not modified; a new model is created
• Weights are randomly initialized (not copied from the original)
• Only compresses models with torch.nn.Linear layers
• Non-linear layers (activations, dropout, batch norm) are preserved
• If the model has fewer than two linear layers, returns a copy

Source code in spectre/core/model.py

def compress(self, ratio: float, preserve_io: bool = True) -> "Model":
    """
    Create a compressed version of the model by reducing layer sizes.

    Creates a new model with reduced hidden layer dimensions while optionally
    preserving input and output dimensions. Only works with sequential models
    containing at least two linear layers.

    Parameters
    ----------
    ratio : float
        Compression ratio between 0 and 1. For example, `ratio=0.5` creates
        a model with half the hidden layer sizes.

    preserve_io : bool, by default True
        Whether to preserve input and output dimensions. If True, only compresses
        hidden layers. If False, compresses all layers except the final output
        layer.

    Returns
    -------
    Model
        New compressed model with reduced layer sizes.

    Raises
    ------
    TypeError
        If ratio is not a float.

    ValueError
        If ratio is not in the range (0, 1).

    Examples
    --------
    >>> model = Model(in_features=100, out_features=10, multipliers=[2.0, 1.5])
    >>> print(model.layer_shapes())
    {'layer_0': (100, 200), 'layer_1': (200, 150), 'layer_2': (150, 10)}
    >>> compressed = model.compress(ratio=0.5, preserve_io=True)
    >>> print(compressed.layer_shapes())
    {'layer_0': (100, 100), 'layer_1': (100, 75), 'layer_2': (75, 10)}

    Compress all layers except output:

    >>> compressed = model.compress(ratio=0.5, preserve_io=False)
    >>> print(compressed.layer_shapes())
    {'layer_0': (50, 100), 'layer_1': (100, 75), 'layer_2': (75, 10)}

    Notes
    -----
    - The original model is not modified; a new model is created
    - Weights are randomly initialized (not copied from the original)
    - Only compresses models with `torch.nn.Linear` layers
    - Non-linear layers (activations, dropout, batch norm) are preserved
    - If the model has fewer than two linear layers, returns a copy
    """
    if not isinstance(ratio, float):
        raise TypeError("Ratio must be a float between 0 and 1.")
    check_in_interval(ratio, "(0, 1)")

    # Get the layers from the sequential
    layers = list(self.sequential)
    linear_indices = [
        i for i, layer in enumerate(layers) if isinstance(layer, torch.nn.Linear)
    ]

    if len(linear_indices) < 2:
        # Not enough layers to compress meaningfully, return a deep copy
        return copy.deepcopy(self)

    # Determine which layers to compress
    if preserve_io:
        # Skip first and last linear layers to preserve input/output dims
        compress_indices = set(linear_indices[1:-1])
    else:
        # Compress all except the last layer (preserve output dim only)
        compress_indices = set(linear_indices[:-1])

    compressed_layers = []
    layer_size_map = {}  # Track size changes for input dimension updates

    for i, layer in enumerate(layers):
        if isinstance(layer, torch.nn.Linear):
            in_features = layer.in_features

            # Update input size if previous layer was compressed
            if i > 0:
                prev_linear_idx = None
                for j in range(i - 1, -1, -1):
                    if isinstance(layers[j], torch.nn.Linear):
                        prev_linear_idx = j
                        break

                if (
                    prev_linear_idx is not None
                    and prev_linear_idx in layer_size_map
                ):
                    in_features = layer_size_map[prev_linear_idx]

            # Determine output size
            if i in compress_indices:
                out_features = max(1, int(layer.out_features * ratio))
                layer_size_map[i] = out_features
            else:
                out_features = layer.out_features
                layer_size_map[i] = out_features

            compressed_layer = torch.nn.Linear(in_features, out_features)

            # Copy bias setting
            if layer.bias is not None:
                compressed_layer.bias = torch.nn.Parameter(
                    torch.zeros(out_features)
                )
            else:
                compressed_layer.bias = None

            compressed_layers.append(compressed_layer)
        else:
            # Keep other layers as-is (activation functions, normalization, etc.)
            compressed_layers.append(copy.deepcopy(layer))

    # Create a new Model instance with the compressed layers
    return Model(layers=compressed_layers)

Functions