Model assembly#

torch_ecg models are assembled bottom-up from registered building blocks. This page explains the layering, the registries that connect the layers, the static shape-inference contract every model honors, and the two output conventions present in the codebase.

The layering#

  1. Building-block configs — torch_ecg.model_configs holds ready-made CFG objects for every backbone family (resnet, mobilenet_v1/v2/v3, vgg16, densenet, multi_scopic, xception, regnet, plus the Transformer attention block and classifier/RNN heads), grouped so a full model config such as ECG_CRNN_CONFIG is just a few references to block configs.

  2. Backbones — feature extractors registered in the BACKBONES registry; any of them can serve as the cnn of a task model without modifying the task model’s code.

  3. Task models — the top level in the MODELS registry (ECG_CRNN, ECG_UNET, …); they own the input/output contract (leads, signal length, classes) and combine a backbone, attention and head into one serializable unit.

Registries#

Registries are the package’s extension points — a name-to-class map with a uniform build(name, **kwargs) entry point:

Registry

Contents

MODELS / BACKBONES

task models / feature-extraction backbones

ATTN_LAYERS

attention blocks (Transformer, SE, …)

OPTIMIZERS / SCHEDULERS

assembled by name from train_config

LOSSES

criterion classes (assembled from train_config.loss)

PREPROCESSORS / AUGMENTERS

preprocessing / online augmentation transforms

Registering a custom backbone requires only a decorator (@BACKBONES.register()); the User Guide shows a worked example.

The compute_output_shape contract#

Every model and backbone implements compute_output_shape(seq_len, batch_size=None), which returns the output shape without running a forward pass. This makes model composition verifiable at construction time: a task model obtains the feature shape from its backbone in order to size its head, and a configuration can be checked without executing it:

from torch_ecg.models import ECG_CRNN
from torch_ecg.model_configs import ECG_CRNN_CONFIG

model = ECG_CRNN(classes=["AF", "I-AVB", "LBBB", "Normal"], n_leads=12,
                 config=ECG_CRNN_CONFIG)
model.compute_output_shape(2000)     # (None, 4)

The classification head globally pools over time, so the logits shape is (batch, n_classes) regardless of input length.

Outputs: tensors and output dataclasses#

The classic task models, such as ECG_CRNN, return plain logit tensors. Newer code — the benchmarks being the reference — wraps model outputs in typed dataclasses from torch_ecg.components.outputs (ClassificationOutput, MultiLabelClassificationOutput, WaveDelineationOutput, RPeaksDetectionOutput, …): fields carry explicit names (logits, bin_pred, pred_classes, …), the objects support + for batching and remain plain CFG dicts underneath, so checkpoints and logging keep working. When writing a new task model, prefer the dataclass contract.

System design#

The overall workflow of the package — data files flow through readers, preprocessing and augmentation into tensors, which feed the trainer driving the model, while the config file drives model construction and checkpoints enable resuming:

Figure made with TikZ

Main Modules and Work Flow.