Layers

The primary trainable preprocessing layer is ConnectedFilterPreprocessingLayer.

Primary CFP Layer

class mtlearn.layers.ConnectedFilterPreprocessingLayer(in_channels, filter_specs, *, device='cpu', scale_mode='dataset_clipped_zscore01', eps=1e-06, score_sharpness=1.0, clamp=None, clipped_zscore_radius=3.0, clipped_zscore_floor=0.05, attribute_dtype=None, tos_interpolation=None, tos_infinity_seed_row=0, tos_infinity_seed_col=0)

Bases: Module

Learnable CFP layer defined by per-output filter specifications.

Each item in filter_specs defines one output per input channel: morphology tree, scoring attributes, and reconstructed altitude signal. clamp optionally bounds score_sharpness * logits before the sigmoid.

Tree construction and attribute computation happen outside autograd. Scoring models provide the trainable node-wise sigmoid gates from normalized attributes.

Parameters:

• scale_mode (str)

• eps (float)

• score_sharpness (float)

• clipped_zscore_radius (float)

• clipped_zscore_floor (float)

• tos_infinity_seed_row (int)

• tos_infinity_seed_col (int)

cached_sample_count()

Return the number of cached sample/channel entries.

Return type:

int

forward(x)

Apply all filter specs and return (B, C * specs, H, W).

Parameters:

x (Tensor) – Input tensor shaped (B, C, H, W) or the cached-loader form (x, idx) produced by build_dataloader_cached.

Return type:

Tensor

Returns:

Tensor with one output channel per input channel and filter spec.

regularization_penalty(x)

Return the per-spec training regularization penalty.

Specs without effective regularizers are skipped. Callers must add this scalar to their training objective explicitly.

Return type:

Tensor

Parameters:

x (torch.Tensor)

predict(x, score_sharpness=1000.0)

Run inference with a caller-provided score sharpness.

The method temporarily switches the module to evaluation mode, runs forward under torch.no_grad(), restores per-spec sharpness, and restores the previous training/eval state.

Return type:

Tensor

Parameters:

• x (torch.Tensor)

• score_sharpness (float)

inspect_training_sample(img, channel=0, idx=None, build_if_missing=True)

Return cached or direct attributes, altitude increments, and parameters per spec.

Parameters:

• img (Tensor) – Image tensor shaped (H, W) or (C, H, W).

• channel (int) – Channel to inspect when img has multiple channels.

• idx (int | None) – Optional stable dataset index used to look up cached payloads.

• build_if_missing (bool) – Build a temporary tree payload when no cache entry exists for idx.

Returns:

Dictionary keyed by filter-spec name. Each entry contains raw and normalized attributes, altitude increments, and current trainable parameters.

freeze_ds_stats()

Stop updating dataset-level normalization statistics.

unfreeze_ds_stats()

Resume updating dataset-level normalization statistics.

refresh_cached_normalization()

Recompute normalized attributes for all cached samples.

save_stats(path)

Save dataset-level normalization statistics.

The payload is a torch-safe dictionary containing a format version, scale_mode, and serialized dataset statistics. Per-sample caches are not saved.

Parameters:

path (str)

load_stats(path, refresh_cache=True)

Load dataset-level normalization statistics.

Parameters:

• path (str) – File previously written by save_stats.

• refresh_cache (bool) – Whether to recompute normalized cached attributes immediately after loading.

get_config()

Return the architecture/configuration needed to reconstruct the layer.

The returned dictionary is serializable and accepted by from_config. It describes layer structure, filter specs, normalization mode, sigmoid gain, clamp bounds, and clipped z-score normalization constants. It does not include trainable weights or dataset statistics.

Return type:

dict[str, Any]

get_parameter_contract()

Return parameter names and shapes owned by this CFP layer.

Return type:

dict[str, Any]

get_inference_contract()

Return the CFP contract that defines forward/inference semantics.

Return type:

dict[str, Any]

get_training_contract()

Return training-only CFP settings such as regularization weights.

Return type:

dict[str, Any]

get_contracts()

Return named CFP contracts for parameters, inference, and training.

Return type:

dict[str, Any]

classmethod from_config(config, *, device=None)

Reconstruct a layer from get_config() output.

Return type:

ConnectedFilterPreprocessingLayer

Parameters:

config (Mapping[str, Any])

get_extra_state()

Embed persistent CFP state in PyTorch checkpoints.

This includes the inference contract and dataset normalization statistics. Per-sample tree/attribute caches are intentionally not persisted.

Return type:

dict[str, Any]

set_extra_state(state)

Restore persistent CFP state from state_dict and validate compatibility.

Return type:

None

Parameters:

state (Any)

export_params(path)

Export CFP parameters and metadata for inspection.

This is not the recommended training checkpoint API. Use mtlearn.layers.save_checkpoint for full PyTorch models.

Parameters:

path (str)

init_identity(p0=0.995, *, strict=True)

Initialize scoring models so CFP outputs start close to identity.

Each scoring model owns its initialization semantics. Linear scorers set weights to zero and choose a bias that yields score ~= p0; nonlinear scorers may use another parameterization with the same score-level contract.

Return type:

tuple[str, ...]

Parameters:

• p0 (float)

• strict (bool)

build_dataloader_cached(dataloader)

Wrap a DataLoader and precompute CFP caches/statistics.

The returned DataLoader yields ((x, idx), y) batches with stable dataset indices. During the prepass, this layer builds tree payloads and updates dataset-level statistics for every sample/channel/tree key. Statistics are frozen and cached normalizations are refreshed before the wrapped loader is returned.

The wrapped dataset must return samples whose first item collates into image tensors shaped (B, C, H, W) with finite, non-negative intensities in uint8, normalized [0, 1], or direct [0, 255] scale. These constraints ensure conversion to uint8 preserves the gray-level ordering used to build morphology trees.

build_dataloader_cached_fixed_stats(dataloader, *, index_offset=0)

Wrap a DataLoader and precompute CFP caches without updating stats.

Use this for validation/test splits after training statistics have been built with build_dataloader_cached(...) or restored with load_stats(...). The returned DataLoader yields ((x, idx + index_offset), y) so callers can keep split cache keys disjoint. The same image-intensity contract as build_dataloader_cached(...) applies.

Parameters:

index_offset (int)

Checkpoint Helpers

mtlearn.layers.collect_cfp_configs(model)

Return serializable configs for every primary CFP layer in model.

Return type:

dict[str, dict[str, Any]]

Parameters:

model (torch.nn.Module)

mtlearn.layers.save_checkpoint(path, model)

Save a PyTorch checkpoint with automatically collected CFP configs.

The model parameters are saved with the normal PyTorch state_dict. Primary CFP layers are discovered by module name and their constructor configs are saved separately so a caller can reconstruct the model before calling load_state_dict. The payload intentionally contains only model weights and CFP configs that define the meaning and shape of CFP-related weights.

Return type:

dict[str, Any]

Parameters:

model (torch.nn.Module)

mtlearn.layers.load_checkpoint(path, model_or_factory, *, device=None, strict=True, weights_only=True)

Load a checkpoint saved by save_checkpoint.

model_or_factory can be either an already constructed module or a callable that returns a module. Factories may accept no arguments when the model constructor already hard-codes its CFP layers, or one positional cfp_configs argument when they need the saved CFP configs.

Return type:

tuple[Module, dict[str, Any]]

Parameters:

• model_or_factory (torch.nn.Module | Callable[[...], torch.nn.Module])

• strict (bool)

• weights_only (bool)