Models Module¶

Neural network architectures for nanopore signal classification.

Overview¶

The models module contains 29 PyTorch model architectures across 5 families. All models accept a signal_in_channels parameter for multi-channel signal input (default: 1).

Architectures come from two places, both reachable via get_model() and MODEL_REGISTRY:

• TOML configs in leech/models/configs/ — the ConvLSTM and TCN families (22 of the 29 names). Each config is a fully declarative kind = "graph" architecture built from the layer registry in leech.models.nn, with one [[variants]] entry per registry name. Adding a normalization/pooling variant is a [[variants]] entry, not a new class.
• Hand-written classes for architectures not yet converted (Remora-compatible ConvLSTM, Transformer, ConvOnly, ResNet, SignalCNN).

Discovery is torch-free, so listing model names (e.g. for CLI --model choices) does not pay the torch import cost.

Model Registry¶

get_model(model_name: str, **kwargs: Any) -> nn.Module

152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170

def get_model(model_name: str, **kwargs: Any) -> nn.Module:
    """
    Get model by name.

    Args:
        model_name: Name of model architecture
        **kwargs: Model-specific parameters (passed to model constructor)

    Returns:
        Instantiated model

    Raises:
        ValueError: If model_name not in registry
    """
    if model_name not in MODEL_REGISTRY:
        available = ", ".join(sorted(MODEL_REGISTRY.keys()))
        raise ValueError(f"Unknown model '{model_name}'. Available: {available}")

    return MODEL_REGISTRY[model_name](**kwargs)

Config-Driven Architectures (ConvLSTM and TCN)¶

The ConvLSTM and TCN families are declared in TOML configs, not Python classes. build_model_class() turns each declaration into a real class on first access, so get_model("ConvLSTMDwell"), isinstance() checks, and checkpoint loading all behave as if the class were hand-written.

ConvLSTMDwell (multi-branch Conv-LSTM with dwell time features) is the recommended default. It has three branches — signal (Conv1d on raw signal), sequence (Conv1d on one-hot k-mers), and features (Conv1d on dwell + level statistics) — merged into a BiLSTM followed by a fully connected head. ConvLSTMBase is the same architecture without the feature branch; compare the two to measure the impact of dwell features.

The TCN family replaces the BiLSTM with stacks of dilated causal convolutions. Residual variants take a 2-channel signal input (raw + k-mer model residual), SplitResidual keeps separate branches for the raw signal and the residual, Motor adds motor-region pooling, and DwellAttn adds dwell-only cross-attention.

Config Loader¶

discover_configs() -> dict[str, tuple[str, str | None]]

build_model_class(name: str) -> type

Layer Registry¶

Serial(sublayers: list[Module])

Stack(sublayers: list[Module])

Parallel(sublayers: dict[str, Module], mode: str = 'cat', dim: int = 1)

Graph(nodes: list[dict], inputs: list[str] | None = None, output: str = 'logits', attrs: dict[str, Any] | None = None, build_order: list[str] | None = None)

register(layer: LayerT) -> LayerT

to_dict(layer: Any, include_weights: bool = False) -> dict

from_dict(model_dict: Any, layer_types: dict[str, type] | None = None) -> Any

Remora-compatible Architectures¶

ConvLSTMRemora(signal_len: int = 400, kmer_len: int = 11, num_features: int = DEFAULT_NUM_FEATURES, size: int = DEFAULT_REMORA_SIZE, num_out: int = 2, dropout: float = DEFAULT_REMORA_DROPOUT, seq_encoding: str = 'signal_kmer', signal_kmer_context: tuple[int, int] = (4, 4), signal_in_channels: int = 1)

forward(signal: Tensor, sequence: Tensor, features: Tensor) -> torch.Tensor

predict_proba(*args, **kwargs) -> torch.Tensor

ConvLSTMRemoraBase(signal_len: int = 400, kmer_len: int = 11, size: int = DEFAULT_REMORA_SIZE, num_out: int = 2, dropout: float = DEFAULT_REMORA_DROPOUT, seq_encoding: str = 'signal_kmer', signal_kmer_context: tuple[int, int] = (4, 4), signal_in_channels: int = 1, **kwargs)

forward(signal: Tensor, sequence: Tensor) -> torch.Tensor

predict_proba(*args, **kwargs) -> torch.Tensor

Transformer Architectures¶

TransformerDwell¶

Transformer with multi-head self-attention and dwell features.

TransformerDwell(signal_len: int = DEFAULT_SIGNAL_LEN, kmer_len: int = DEFAULT_KMER_LEN, num_features: int = DEFAULT_NUM_FEATURES, dwell_margin: int = DEFAULT_DWELL_MARGIN, d_model: int = 256, nhead: int = 8, num_layers: int = 4, dim_feedforward: int = 1024, dropout: float = DEFAULT_DROPOUT, seq_encoding: str = 'base_onehot', signal_kmer_context: tuple[int, int] = DEFAULT_SIGNAL_KMER_CONTEXT, signal_in_channels: int = 1, **kwargs)

forward(signal: Tensor, sequence: Tensor, features: Tensor) -> torch.Tensor

TransformerDwellResidual¶

Transformer with 2-channel signal input (raw + kmer residual).

TransformerDwellResidual(*, signal_in_channels: int = 2, **kwargs)

Convolutional Architectures¶

ConvOnly¶

Pure convolutional network with multi-scale convolutions.

ConvOnly(signal_len: int = DEFAULT_SIGNAL_LEN, kmer_len: int = DEFAULT_KMER_LEN, num_features: int = DEFAULT_NUM_FEATURES, dwell_margin: int = DEFAULT_DWELL_MARGIN, base_channels: int = 16, num_blocks: int = 3, num_attn_heads: int = 4, dropout: float = DEFAULT_DROPOUT, seq_encoding: str = 'base_onehot', signal_kmer_context: tuple[int, int] = DEFAULT_SIGNAL_KMER_CONTEXT, signal_in_channels: int = 1, **kwargs)

forward(signal: Tensor, sequence: Tensor, features: Tensor) -> torch.Tensor

ResNetDwell¶

Deep residual network with skip connections.

ResNetDwell(signal_len: int = DEFAULT_SIGNAL_LEN, kmer_len: int = DEFAULT_KMER_LEN, num_features: int = DEFAULT_NUM_FEATURES, dwell_margin: int = DEFAULT_DWELL_MARGIN, base_channels: int = 64, num_attn_heads: int = 4, dropout: float = DEFAULT_DROPOUT, seq_encoding: str = 'base_onehot', signal_kmer_context: tuple[int, int] = DEFAULT_SIGNAL_KMER_CONTEXT, signal_in_channels: int = 1, **kwargs)

forward(signal: Tensor, sequence: Tensor, features: Tensor) -> torch.Tensor

SignalCNN¶

Signal-only 1D-CNN classifier (ignores sequence and dwell inputs).

SignalCNN(num_classes: int = 2, signal_len: int = 256, channels: int = 32, signal_in_channels: int = 1, kernel_size: int = 7, seq_encoding: str = 'base_onehot', **_: object)

Inference Wrappers¶

ModelInferenceWrapper(model: Module, model_type: str)

parameters

enable_repr_capture() -> int

forward_batch(batch: dict, device: str) -> torch.Tensor

train() -> None

eval() -> None

to(device: str) -> ModelInferenceWrapper

state_dict()

load_state_dict(state_dict)

RemoraModelWrapper(model_path: str | Path, device: str = 'cpu')

forward_batch(batch: dict, device: str) -> torch.Tensor