deepecgkit.models

Neural network architectures for ECG signal processing. All models are nn.Module subclasses with a consistent interface: __init__(input_channels, output_size, **kwargs).

Base Class

ECGModelBase

Bases: PretrainedMixin, Module

Base class for registered ECG models.

Subclasses must:

• call super().__init__() first,
• set self._feature_dim to the backbone's output width,
• create their head module under the attribute named by HEAD_ATTR,
• implement extract_features.

A subclass without its own top-level self.dropout inherits the parameter-free nn.Identity set here, keeping the default forward (features -> dropout -> head) bit-identical to a model with no dropout.

Source code in deepecgkit/models/base.py

class ECGModelBase(PretrainedMixin, nn.Module):
    """Base class for registered ECG models.

    Subclasses must:

    - call ``super().__init__()`` first,
    - set ``self._feature_dim`` to the backbone's output width,
    - create their head module under the attribute named by ``HEAD_ATTR``,
    - implement ``extract_features``.

    A subclass without its own top-level ``self.dropout`` inherits the
    parameter-free ``nn.Identity`` set here, keeping the default ``forward``
    (features -> dropout -> head) bit-identical to a model with no dropout.
    """

    #: Name of the attribute holding the classification head submodule.
    HEAD_ATTR: ClassVar[str] = "classifier"

    def __init__(self) -> None:
        super().__init__()
        self.dropout: nn.Module = nn.Identity()

    @property
    def feature_dim(self) -> int:
        """Width of the flat feature vector produced by ``extract_features``."""
        return self._feature_dim

    def get_head(self) -> nn.Module:
        """Return the classification head, whatever its attribute name."""
        return getattr(self, self.HEAD_ATTR)

    def set_head(self, module: nn.Module) -> None:
        """Replace the classification head, preserving its attribute name."""
        setattr(self, self.HEAD_ATTR, module)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        """Map an input batch to flat feature vectors of size ``feature_dim``."""
        raise NotImplementedError

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.extract_features(x)
        x = self.dropout(x)
        return self.get_head()(x)

feature_dim property

feature_dim: int

Width of the flat feature vector produced by extract_features.

get_head

get_head() -> nn.Module

Return the classification head, whatever its attribute name.

Source code in deepecgkit/models/base.py

def get_head(self) -> nn.Module:
    """Return the classification head, whatever its attribute name."""
    return getattr(self, self.HEAD_ATTR)

set_head

set_head(module: Module) -> None

Replace the classification head, preserving its attribute name.

Source code in deepecgkit/models/base.py

def set_head(self, module: nn.Module) -> None:
    """Replace the classification head, preserving its attribute name."""
    setattr(self, self.HEAD_ATTR, module)

extract_features

extract_features(x: Tensor) -> torch.Tensor

Map an input batch to flat feature vectors of size feature_dim.

Source code in deepecgkit/models/base.py

def extract_features(self, x: torch.Tensor) -> torch.Tensor:
    """Map an input batch to flat feature vectors of size ``feature_dim``."""
    raise NotImplementedError

Building and Loading

build_model

build_model(
    name: str,
    input_channels: int,
    output_size: int,
    **kwargs: Any,
) -> nn.Module

Instantiate a registered model.

Parameters:

Name	Type	Description	Default
name	str	Registered model name	required
input_channels	int	Number of ECG leads the model consumes	required
output_size	int	Number of output units of the classifier head	required
**kwargs	Any	Additional keyword arguments forwarded to the model, overriding any default_kwargs recorded in the registry	{}

Returns:

Type	Description
Module	The instantiated model.

Raises:

Type	Description
KeyError	If the model name is not registered.

Source code in deepecgkit/models/factory.py

def build_model(name: str, input_channels: int, output_size: int, **kwargs: Any) -> nn.Module:
    """Instantiate a registered model.

    Args:
        name: Registered model name
        input_channels: Number of ECG leads the model consumes
        output_size: Number of output units of the classifier head
        **kwargs: Additional keyword arguments forwarded to the model, overriding
            any default_kwargs recorded in the registry

    Returns:
        The instantiated model.

    Raises:
        KeyError: If the model name is not registered.
    """
    if name not in get_model_names():
        raise KeyError(f"Unknown model '{name}'. Available: {', '.join(get_model_names())}")
    model_class = get_model(name)
    resolved = dict(get_model_info(name).get("default_kwargs") or {})
    resolved.update(kwargs)
    return model_class(input_channels=input_channels, output_size=output_size, **resolved)

build_model_from_checkpoint

build_model_from_checkpoint(
    checkpoint_path: str | Path,
    name: str | None = None,
    input_channels: int | None = None,
    output_size: int | None = None,
    map_location: str = "cpu",
) -> tuple[nn.Module, dict[str, Any]]

Build a model matching a checkpoint and load its weights.

The architecture and its shapes are taken from explicit arguments when given, then from the model_config block recorded at save time, then (for shapes) inferred from the weights themselves.

Checkpoints written by ECGTrainer, bare state dicts, and legacy PyTorch Lightning checkpoints are all accepted.

Parameters:

Name	Type	Description	Default
checkpoint_path	str | Path	Path to a checkpoint written by ECGTrainer.save_checkpoint	required
name	str | None	Registered model name; optional when the checkpoint records one	None
input_channels	int | None	Overrides the recorded or inferred lead count	None
output_size	int | None	Overrides the recorded or inferred output size	None
map_location	str	Device to map tensors onto while loading	'cpu'

Returns:

Type	Description
tuple[Module, dict[str, Any]]	Tuple of the loaded model and the checkpoint dictionary.

Raises:

Type	Description
ValueError	If the checkpoint holds no weights, names a different architecture than the one requested, does not say which architecture it holds, or the lead count cannot be determined.

Source code in deepecgkit/models/factory.py

def build_model_from_checkpoint(
    checkpoint_path: str | Path,
    name: str | None = None,
    input_channels: int | None = None,
    output_size: int | None = None,
    map_location: str = "cpu",
) -> tuple[nn.Module, dict[str, Any]]:
    """Build a model matching a checkpoint and load its weights.

    The architecture and its shapes are taken from explicit arguments when
    given, then from the model_config block recorded at save time, then (for
    shapes) inferred from the weights themselves.

    Checkpoints written by ECGTrainer, bare state dicts, and legacy PyTorch
    Lightning checkpoints are all accepted.

    Args:
        checkpoint_path: Path to a checkpoint written by ECGTrainer.save_checkpoint
        name: Registered model name; optional when the checkpoint records one
        input_channels: Overrides the recorded or inferred lead count
        output_size: Overrides the recorded or inferred output size
        map_location: Device to map tensors onto while loading

    Returns:
        Tuple of the loaded model and the checkpoint dictionary.

    Raises:
        ValueError: If the checkpoint holds no weights, names a different
            architecture than the one requested, does not say which
            architecture it holds, or the lead count cannot be determined.
    """
    checkpoint = load_checkpoint_file(checkpoint_path, map_location)
    try:
        state_dict, recorded = _split_checkpoint(checkpoint)
    except ValueError as error:
        raise ValueError(f"Checkpoint {checkpoint_path} contains no model_state_dict") from error

    recorded_name = recorded_model_name(checkpoint)
    if name is None:
        if recorded_name is None:
            raise ValueError(
                f"Checkpoint {checkpoint_path} does not record which architecture it holds. "
                f"Name it explicitly (CLI: --model)."
            )
        name = recorded_name
    elif recorded_name is not None and recorded_name != name:
        raise ValueError(
            f"Checkpoint {checkpoint_path} was trained as '{recorded_name}', not '{name}'. "
            f"Leave the model name out, or pass '{recorded_name}'."
        )

    if output_size is None:
        head_attr = get_model_info(name).get("head_attr") if name in get_model_names() else None
        output_size = recorded.get("output_size") or infer_output_size(state_dict, head_attr)
    if output_size is None:
        raise ValueError(
            f"Could not determine output size for '{name}' from {checkpoint_path}. "
            "Pass output_size explicitly."
        )

    if input_channels is None:
        input_channels = recorded.get("input_channels") or infer_input_channels(state_dict)
    if input_channels is None:
        raise ValueError(
            f"Input channel count is not recoverable from the weights of '{name}'. "
            "Pass input_channels explicitly, for example from the dataset being evaluated."
        )

    model = build_model(name, input_channels=input_channels, output_size=output_size)
    try:
        model.load_state_dict(state_dict)
    except RuntimeError as error:
        expected = set(model.state_dict())
        missing = len(expected - set(state_dict))
        unexpected = len(set(state_dict) - expected)
        if not missing and not unexpected:
            raise
        raise ValueError(
            f"The weights in {checkpoint_path} do not fit the '{name}' architecture "
            f"({missing} of its tensors are missing and {unexpected} are not part of it). "
            f"The checkpoint was most likely trained with a different model."
        ) from error
    return model, checkpoint

CNN Architectures

SimpleCNN

Bases: ECGModelBase

Simple CNN model for ECG signal classification.

A straightforward convolutional neural network with pooling and dropout layers for basic ECG classification tasks.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
dropout_rate	float	Dropout probability (default: 0.3)	0.3

Example

model = SimpleCNN(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 256)

Source code in deepecgkit/models/simple_cnn.py

@register_model(
    name="simple-cnn",
    description="Lightweight CNN for fast inference",
    min_input_length=8,
)
class SimpleCNN(ECGModelBase):
    """
    Simple CNN model for ECG signal classification.

    A straightforward convolutional neural network with pooling and dropout
    layers for basic ECG classification tasks.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        dropout_rate: Dropout probability (default: 0.3)

    Example:
        >>> model = SimpleCNN(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 256)
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        self.features = nn.Sequential(
            nn.Conv1d(input_channels, 32, kernel_size=7, padding=3),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2),
            nn.Dropout(dropout_rate),
            nn.Conv1d(32, 64, kernel_size=5, padding=2),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2),
            nn.Dropout(dropout_rate),
            nn.Conv1d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.MaxPool1d(kernel_size=2),
            nn.Dropout(dropout_rate),
            nn.Conv1d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.AdaptiveAvgPool1d(1),
        )

        self._feature_dim = 256
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(128, output_size),
        )

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        return torch.flatten(x, 1)

FCNWang

Bases: ECGModelBase

Fully Convolutional Network for ECG signal classification.

Based on Wang et al.'s FCN architecture for time series classification, using three convolutional blocks with batch normalization and global average pooling. No dense layers except the final classifier.

Reference

Wang Z., Yan W., Oates T. "Time Series Classification from Scratch with Deep Neural Networks: A Strong Baseline" (2017) https://github.com/helme/ecg_ptbxl_benchmarking

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 12 for 12-lead ECG)	12
output_size	int	Number of output classes	5
filters	list[int] | None	List of filter counts for each conv block (default: [128, 256, 128])	None
kernel_sizes	list[int] | None	List of kernel sizes for each conv block (default: [8, 5, 3])	None
dropout_rate	float	Dropout probability before classifier (default: 0.3)	0.3

Example

model = FCNWang(input_channels=12, output_size=5) x = torch.randn(32, 12, 1000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 128)

Source code in deepecgkit/models/fcn_wang.py

@register_model(
    name="fcn-wang",
    description="Fully Convolutional Network (Wang et al.) for time series classification",
    min_input_length=2,
)
class FCNWang(ECGModelBase):
    """
    Fully Convolutional Network for ECG signal classification.

    Based on Wang et al.'s FCN architecture for time series classification,
    using three convolutional blocks with batch normalization and global
    average pooling. No dense layers except the final classifier.

    Reference:
        Wang Z., Yan W., Oates T. "Time Series Classification from Scratch
        with Deep Neural Networks: A Strong Baseline" (2017)
        https://github.com/helme/ecg_ptbxl_benchmarking

    Args:
        input_channels: Number of input channels (default: 12 for 12-lead ECG)
        output_size: Number of output classes
        filters: List of filter counts for each conv block (default: [128, 256, 128])
        kernel_sizes: List of kernel sizes for each conv block (default: [8, 5, 3])
        dropout_rate: Dropout probability before classifier (default: 0.3)

    Example:
        >>> model = FCNWang(input_channels=12, output_size=5)
        >>> x = torch.randn(32, 12, 1000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 128)
    """

    def __init__(
        self,
        input_channels: int = 12,
        output_size: int = 5,
        filters: list[int] | None = None,
        kernel_sizes: list[int] | None = None,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        if filters is None:
            filters = [128, 256, 128]
        if kernel_sizes is None:
            kernel_sizes = [8, 5, 3]

        if len(filters) != len(kernel_sizes):
            raise ValueError(
                f"filters and kernel_sizes must have the same length, got "
                f"{len(filters)} and {len(kernel_sizes)}"
            )

        blocks = []
        in_ch = input_channels
        for nf, ks in zip(filters, kernel_sizes, strict=False):
            blocks.append(
                nn.Sequential(
                    nn.Conv1d(in_ch, nf, kernel_size=ks, padding=(ks - 1) // 2, bias=False),
                    nn.BatchNorm1d(nf),
                    nn.ReLU(inplace=True),
                )
            )
            in_ch = nf

        self.conv_blocks = nn.Sequential(*blocks)
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = filters[-1]
        self.classifier = nn.Linear(self._feature_dim, output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv_blocks(x)
        x = self.global_pool(x)
        return torch.flatten(x, 1)

AFModel

Bases: ECGModelBase

Atrial Fibrillation classification model.

A convolutional neural network specifically designed for AF detection in ECG signals with configurable recording length support.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input ECG leads (default: 1)	1
output_size	int	Number of output classes (default: 4)	4
recording_length	int	Recording length in seconds (must be 6, 10, or 30)	30

Example

model = AFModel(recording_length=30) x = torch.randn(32, 1, 9000) output = model(x) print(output.shape) # (32, 4)

features = model.extract_features(x) print(features.shape) # (32, 196)

Source code in deepecgkit/models/af_classifier.py

@register_model(
    name="afmodel",
    description="Atrial Fibrillation model optimized for 30s ECG segments",
    default_kwargs={"recording_length": 30},
    min_input_length=7971,
)
class AFModel(ECGModelBase):
    """
    Atrial Fibrillation classification model.

    A convolutional neural network specifically designed for AF detection
    in ECG signals with configurable recording length support.

    Args:
        input_channels: Number of input ECG leads (default: 1)
        output_size: Number of output classes (default: 4)
        recording_length: Recording length in seconds (must be 6, 10, or 30)

    Example:
        >>> model = AFModel(recording_length=30)
        >>> x = torch.randn(32, 1, 9000)
        >>> output = model(x)
        >>> print(output.shape)  # (32, 4)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 196)
    """

    HEAD_ATTR: ClassVar[str] = "output"

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        recording_length: int = 30,
    ):
        super().__init__()
        if recording_length not in (6, 10, 30):
            raise ValueError(f"Recording length must be 6, 10, or 30, got {recording_length}")

        self.recording_length = recording_length
        kernel_size = 6 if recording_length in [6, 10] else 32

        layer_configs = [
            (input_channels, 64, True, True),
            (64, 64, False, False),
            (64, 64, True, True),
            (64, 64, False, False),
            (64, 64, True, True),
            (64, 64, False, False),
            (64, 128, True, True),
            (128, 128, False, False),
            (128, 128, True, True),
            (128, 128, False, False),
            (128, 196, True, True),
            (196, 196, False, False),
            (196, 196, True, True),
        ]

        self.conv_layers = nn.ModuleList(
            [
                ConvBlock(in_ch, out_ch, kernel_size, dropout, pooling)
                for in_ch, out_ch, dropout, pooling in layer_configs
            ]
        )

        self.adaptive_pool = nn.AdaptiveAvgPool1d(1)
        self.flatten = nn.Flatten()
        self._feature_dim = 196
        self.dropout = nn.Dropout(0.3)
        self.output = nn.Linear(in_features=self._feature_dim, out_features=output_size)

    @property
    def min_input_length(self) -> int:
        """Smallest signal length the convolutional stack accepts.

        Each block narrows the signal by ``kernel_size - 1`` and optionally
        halves it, so the bound is obtained by walking the stack backwards from
        a single remaining sample.

        Returns:
            Minimum number of samples per lead.
        """
        required = 1
        for block in reversed(self.conv_layers):
            if block.use_pooling:
                required *= 2
            required += block.conv.kernel_size[0] - 1
        return required

    def extract_features(self, data: torch.Tensor) -> torch.Tensor:
        """Return the pooled representation preceding dropout and the head.

        Args:
            data: Input of shape (batch, leads, time)

        Returns:
            Tensor of shape (batch, feature_dim).

        Raises:
            ValueError: If the signal is shorter than :attr:`min_input_length`.
                ``recording_length`` selects the convolution kernel size but is
                never compared against the data, so without this check a short
                signal fails deep inside the stack with an opaque kernel-size
                error.
        """
        if data.dim() != 3:
            raise ValueError(
                f"{type(self).__name__} expects input of shape (batch, leads, time), "
                f"got a {data.dim()}-D tensor of shape {tuple(data.shape)}."
            )
        if data.shape[-1] < self.min_input_length:
            raise ValueError(
                f"{type(self).__name__}(recording_length={self.recording_length}) needs at "
                f"least {self.min_input_length} samples per lead but received "
                f"{data.shape[-1]}. Pass longer signals, or build the model with a "
                f"recording_length matching the data."
            )

        for conv_layer in self.conv_layers:
            data = conv_layer(data)
        data = self.adaptive_pool(data)
        data = self.flatten(data)
        return data

min_input_length property

min_input_length: int

Smallest signal length the convolutional stack accepts.

Each block narrows the signal by kernel_size - 1 and optionally halves it, so the bound is obtained by walking the stack backwards from a single remaining sample.

Returns:

Type	Description
int	Minimum number of samples per lead.

extract_features

extract_features(data: Tensor) -> torch.Tensor

Return the pooled representation preceding dropout and the head.

Parameters:

Name	Type	Description	Default
data	Tensor	Input of shape (batch, leads, time)	required

Returns:

Type	Description
Tensor	Tensor of shape (batch, feature_dim).

Raises:

Type	Description
ValueError	If the signal is shorter than :attr:min_input_length. recording_length selects the convolution kernel size but is never compared against the data, so without this check a short signal fails deep inside the stack with an opaque kernel-size error.

Source code in deepecgkit/models/af_classifier.py

def extract_features(self, data: torch.Tensor) -> torch.Tensor:
    """Return the pooled representation preceding dropout and the head.

    Args:
        data: Input of shape (batch, leads, time)

    Returns:
        Tensor of shape (batch, feature_dim).

    Raises:
        ValueError: If the signal is shorter than :attr:`min_input_length`.
            ``recording_length`` selects the convolution kernel size but is
            never compared against the data, so without this check a short
            signal fails deep inside the stack with an opaque kernel-size
            error.
    """
    if data.dim() != 3:
        raise ValueError(
            f"{type(self).__name__} expects input of shape (batch, leads, time), "
            f"got a {data.dim()}-D tensor of shape {tuple(data.shape)}."
        )
    if data.shape[-1] < self.min_input_length:
        raise ValueError(
            f"{type(self).__name__}(recording_length={self.recording_length}) needs at "
            f"least {self.min_input_length} samples per lead but received "
            f"{data.shape[-1]}. Pass longer signals, or build the model with a "
            f"recording_length matching the data."
        )

    for conv_layer in self.conv_layers:
        data = conv_layer(data)
    data = self.adaptive_pool(data)
    data = self.flatten(data)
    return data

Residual Networks

ResNet1D

Bases: ECGModelBase

1D ResNet model for ECG signal classification.

A residual network architecture adapted for 1D time-series ECG data, providing deep feature extraction with skip connections.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
base_channels	int	Base number of channels (default: 64)	64
num_blocks	list | None	List of number of blocks in each layer (default: [2, 2, 2, 2])	None
dropout_rate	float	Dropout probability (default: 0.3)	0.3

Example

model = ResNet1D(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 512) with base_channels=64

Source code in deepecgkit/models/resnet1d.py

@register_model(
    name="resnet",
    description="1D ResNet architecture adapted for ECG",
)
class ResNet1D(ECGModelBase):
    """
    1D ResNet model for ECG signal classification.

    A residual network architecture adapted for 1D time-series ECG data,
    providing deep feature extraction with skip connections.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        base_channels: Base number of channels (default: 64)
        num_blocks: List of number of blocks in each layer (default: [2, 2, 2, 2])
        dropout_rate: Dropout probability (default: 0.3)

    Example:
        >>> model = ResNet1D(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 512) with base_channels=64
    """

    HEAD_ATTR: ClassVar[str] = "fc"

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        base_channels: int = 64,
        num_blocks: list | None = None,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        if num_blocks is None:
            num_blocks = [2, 2, 2, 2]

        self.in_channels = base_channels

        self.conv1 = nn.Conv1d(
            input_channels,
            base_channels,
            kernel_size=15,
            stride=2,
            padding=7,
            bias=False,
        )
        self.bn1 = nn.BatchNorm1d(base_channels)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)

        self.layer1 = self._make_layer(base_channels, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(base_channels * 2, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(base_channels * 4, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(base_channels * 8, num_blocks[3], stride=2)

        self.avgpool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = base_channels * 8
        self.fc = nn.Linear(self._feature_dim, output_size)

    def _make_layer(self, out_channels: int, num_blocks: int, stride: int = 1) -> nn.Sequential:
        downsample = None
        if stride != 1 or self.in_channels != out_channels:
            downsample = nn.Sequential(
                nn.Conv1d(
                    self.in_channels,
                    out_channels,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm1d(out_channels),
            )

        layers = []
        layers.append(
            ResidualBlock1D(
                self.in_channels,
                out_channels,
                stride=stride,
                downsample=downsample,
            )
        )
        self.in_channels = out_channels

        for _ in range(1, num_blocks):
            layers.append(ResidualBlock1D(out_channels, out_channels))

        return nn.Sequential(*layers)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return x

ResNetWang

Bases: ECGModelBase

ResNet model based on Wang et al.'s architecture for time series classification.

A shallow 3-block residual network without initial pooling, using larger initial channels (128) and asymmetric kernel sizes in residual blocks. This is the standard ResNet baseline from the PTB-XL benchmark.

Reference

Wang Z., Yan W., Oates T. "Time Series Classification from Scratch with Deep Neural Networks: A Strong Baseline" (2017) https://github.com/helme/ecg_ptbxl_benchmarking

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 12 for 12-lead ECG)	12
output_size	int	Number of output classes	5
base_channels	int	Base number of channels (default: 128)	128
kernel_size	int	Primary kernel size for residual blocks (default: 5)	5
kernel_size_stem	int	Kernel size for the stem convolution (default: 7)	7
dropout_rate	float	Dropout probability before classifier (default: 0.3)	0.3

Example

model = ResNetWang(input_channels=12, output_size=5) x = torch.randn(32, 12, 1000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 128)

Source code in deepecgkit/models/resnet1d_wang.py

@register_model(
    name="resnet-wang",
    description="ResNet (Wang et al.) for time series classification",
)
class ResNetWang(ECGModelBase):
    """
    ResNet model based on Wang et al.'s architecture for time series classification.

    A shallow 3-block residual network without initial pooling, using larger
    initial channels (128) and asymmetric kernel sizes in residual blocks.
    This is the standard ResNet baseline from the PTB-XL benchmark.

    Reference:
        Wang Z., Yan W., Oates T. "Time Series Classification from Scratch
        with Deep Neural Networks: A Strong Baseline" (2017)
        https://github.com/helme/ecg_ptbxl_benchmarking

    Args:
        input_channels: Number of input channels (default: 12 for 12-lead ECG)
        output_size: Number of output classes
        base_channels: Base number of channels (default: 128)
        kernel_size: Primary kernel size for residual blocks (default: 5)
        kernel_size_stem: Kernel size for the stem convolution (default: 7)
        dropout_rate: Dropout probability before classifier (default: 0.3)

    Example:
        >>> model = ResNetWang(input_channels=12, output_size=5)
        >>> x = torch.randn(32, 12, 1000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 128)
    """

    def __init__(
        self,
        input_channels: int = 12,
        output_size: int = 5,
        base_channels: int = 128,
        kernel_size: int = 5,
        kernel_size_stem: int = 7,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        self.in_channels = base_channels

        self.stem = nn.Sequential(
            nn.Conv1d(
                input_channels,
                base_channels,
                kernel_size=kernel_size_stem,
                stride=1,
                padding=(kernel_size_stem - 1) // 2,
                bias=False,
            ),
            nn.BatchNorm1d(base_channels),
            nn.ReLU(inplace=True),
        )

        self.layer1 = self._make_layer(base_channels, 1, kernel_size=kernel_size, stride=1)
        self.layer2 = self._make_layer(base_channels, 1, kernel_size=kernel_size, stride=1)
        self.layer3 = self._make_layer(base_channels, 1, kernel_size=kernel_size, stride=1)

        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = base_channels
        self.classifier = nn.Linear(self._feature_dim, output_size)

    def _make_layer(
        self,
        out_channels: int,
        num_blocks: int,
        kernel_size: int = 5,
        stride: int = 1,
    ) -> nn.Sequential:
        downsample = None
        if stride != 1 or self.in_channels != out_channels:
            downsample = nn.Sequential(
                nn.Conv1d(self.in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm1d(out_channels),
            )

        layers = [
            BasicBlock1DWang(
                self.in_channels,
                out_channels,
                kernel_size=kernel_size,
                stride=stride,
                downsample=downsample,
            )
        ]
        self.in_channels = out_channels
        for _ in range(1, num_blocks):
            layers.append(BasicBlock1DWang(out_channels, out_channels, kernel_size=kernel_size))

        return nn.Sequential(*layers)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.global_pool(x)
        return torch.flatten(x, 1)

SEResNet1D

Bases: ECGModelBase

SE-ResNet model for ECG signal classification.

Extends ResNet1D with Squeeze-and-Excitation blocks that learn channel attention weights. Particularly effective for multi-lead ECG where the network can learn which leads are most informative for each class.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
base_channels	int	Base number of channels (default: 64)	64
num_blocks	list | None	List of number of blocks in each layer (default: [2, 2, 2, 2])	None
se_reduction	int	SE reduction ratio (default: 16)	16
dropout_rate	float	Dropout probability (default: 0.3)	0.3

Example

model = SEResNet1D(input_channels=12, output_size=5) x = torch.randn(32, 12, 5000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 512) with base_channels=64

Source code in deepecgkit/models/se_resnet1d.py

@register_model(
    name="se-resnet",
    description="ResNet with Squeeze-and-Excitation channel attention for ECG",
)
class SEResNet1D(ECGModelBase):
    """
    SE-ResNet model for ECG signal classification.

    Extends ResNet1D with Squeeze-and-Excitation blocks that learn channel
    attention weights. Particularly effective for multi-lead ECG where the
    network can learn which leads are most informative for each class.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        base_channels: Base number of channels (default: 64)
        num_blocks: List of number of blocks in each layer (default: [2, 2, 2, 2])
        se_reduction: SE reduction ratio (default: 16)
        dropout_rate: Dropout probability (default: 0.3)

    Example:
        >>> model = SEResNet1D(input_channels=12, output_size=5)
        >>> x = torch.randn(32, 12, 5000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 512) with base_channels=64
    """

    HEAD_ATTR: ClassVar[str] = "fc"

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        base_channels: int = 64,
        num_blocks: list | None = None,
        se_reduction: int = 16,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        if num_blocks is None:
            num_blocks = [2, 2, 2, 2]

        self.in_channels = base_channels
        self.se_reduction = se_reduction
        self.dropout_rate = dropout_rate

        self.conv1 = nn.Conv1d(
            input_channels,
            base_channels,
            kernel_size=15,
            stride=2,
            padding=7,
            bias=False,
        )
        self.bn1 = nn.BatchNorm1d(base_channels)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)

        self.layer1 = self._make_layer(base_channels, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(base_channels * 2, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(base_channels * 4, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(base_channels * 8, num_blocks[3], stride=2)

        self.avgpool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = base_channels * 8
        self.fc = nn.Linear(self._feature_dim, output_size)

    def _make_layer(self, out_channels: int, num_blocks: int, stride: int = 1) -> nn.Sequential:
        downsample = None
        if stride != 1 or self.in_channels != out_channels:
            downsample = nn.Sequential(
                nn.Conv1d(
                    self.in_channels,
                    out_channels,
                    kernel_size=1,
                    stride=stride,
                    bias=False,
                ),
                nn.BatchNorm1d(out_channels),
            )

        layers = [
            SEResidualBlock1D(
                self.in_channels,
                out_channels,
                stride=stride,
                downsample=downsample,
                se_reduction=self.se_reduction,
                dropout_rate=self.dropout_rate,
            )
        ]
        self.in_channels = out_channels

        for _ in range(1, num_blocks):
            layers.append(
                SEResidualBlock1D(
                    out_channels,
                    out_channels,
                    se_reduction=self.se_reduction,
                    dropout_rate=self.dropout_rate,
                )
            )

        return nn.Sequential(*layers)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return x

XResNet1D

Bases: ECGModelBase

XResNet model for ECG signal classification.

An improved ResNet incorporating three key enhancements from recent research: (1) a multi-layer stem instead of a single large convolution, (2) Mish activation for smoother gradients, and (3) anti-aliased blur-pool downsampling to reduce aliasing.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
base_channels	int	Base number of channels (default: 64)	64
num_blocks	list | None	List of number of blocks in each layer (default: [2, 2, 2, 2])	None
dropout_rate	float	Dropout probability (default: 0.3)	0.3
use_blur_pool	bool	Whether to use anti-aliased downsampling (default: True)	True

Example

model = XResNet1D(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 512) with base_channels=64

Source code in deepecgkit/models/xresnet1d.py

@register_model(
    name="xresnet",
    description="Improved ResNet with Mish activation and blur-pool downsampling",
    min_input_length=17,
)
class XResNet1D(ECGModelBase):
    """
    XResNet model for ECG signal classification.

    An improved ResNet incorporating three key enhancements from recent
    research: (1) a multi-layer stem instead of a single large convolution,
    (2) Mish activation for smoother gradients, and (3) anti-aliased
    blur-pool downsampling to reduce aliasing.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        base_channels: Base number of channels (default: 64)
        num_blocks: List of number of blocks in each layer (default: [2, 2, 2, 2])
        dropout_rate: Dropout probability (default: 0.3)
        use_blur_pool: Whether to use anti-aliased downsampling (default: True)

    Example:
        >>> model = XResNet1D(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 512) with base_channels=64
    """

    HEAD_ATTR: ClassVar[str] = "fc"

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        base_channels: int = 64,
        num_blocks: list | None = None,
        dropout_rate: float = 0.3,
        use_blur_pool: bool = True,
    ):
        super().__init__()

        if num_blocks is None:
            num_blocks = [2, 2, 2, 2]

        self.in_channels = base_channels
        self.use_blur_pool = use_blur_pool
        self.act = Mish()

        self.stem = nn.Sequential(
            nn.Conv1d(
                input_channels,
                base_channels // 2,
                kernel_size=7,
                stride=2,
                padding=3,
                bias=False,
            ),
            nn.BatchNorm1d(base_channels // 2),
            Mish(),
            nn.Conv1d(
                base_channels // 2,
                base_channels // 2,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False,
            ),
            nn.BatchNorm1d(base_channels // 2),
            Mish(),
            nn.Conv1d(
                base_channels // 2,
                base_channels,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False,
            ),
            nn.BatchNorm1d(base_channels),
            Mish(),
        )

        if use_blur_pool:
            self.stem_pool = BlurPool1D(base_channels, stride=2)
        else:
            self.stem_pool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)

        self.layer1 = self._make_layer(base_channels, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(base_channels * 2, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(base_channels * 4, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(base_channels * 8, num_blocks[3], stride=2)

        self.avgpool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = base_channels * 8
        self.fc = nn.Linear(self._feature_dim, output_size)

    def _make_layer(self, out_channels: int, num_blocks: int, stride: int = 1) -> nn.Sequential:
        downsample = None
        if stride != 1 or self.in_channels != out_channels:
            ds_layers: list[nn.Module] = []
            if stride > 1:
                if self.use_blur_pool:
                    ds_layers.append(BlurPool1D(self.in_channels, stride=stride))
                else:
                    ds_layers.append(nn.AvgPool1d(kernel_size=stride, stride=stride))
            ds_layers.append(nn.Conv1d(self.in_channels, out_channels, kernel_size=1, bias=False))
            ds_layers.append(nn.BatchNorm1d(out_channels))
            downsample = nn.Sequential(*ds_layers)

        layers = [
            XResBlock1D(
                self.in_channels,
                out_channels,
                stride=stride,
                downsample=downsample,
                use_blur_pool=self.use_blur_pool,
            )
        ]
        self.in_channels = out_channels

        for _ in range(1, num_blocks):
            layers.append(
                XResBlock1D(
                    out_channels,
                    out_channels,
                    use_blur_pool=self.use_blur_pool,
                )
            )

        return nn.Sequential(*layers)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        x = self.stem_pool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        return x

XResNet1dBenchmark

Bases: ECGModelBase

XResNet1d adapted from the PTB-XL benchmarking repository.

Key differences from the standard XResNet1D: - Fixed feature dimension across all layers (all blocks use same width) - Concat pooling head (AdaptiveAvgPool + AdaptiveMaxPool concatenated) - Multi-layer stem with configurable kernel size - Supports both BasicBlock (expansion=1) and Bottleneck (expansion=4)

Reference

Strodthoff N., Wagner P., Schaeffter T., Samek W. "Deep Learning for ECG Analysis: Benchmarks and Insights from PTB-XL" (2021) https://github.com/helme/ecg_ptbxl_benchmarking

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 12 for 12-lead ECG)	12
output_size	int	Number of output classes	5
expansion	int	Block expansion factor (1=BasicBlock, 4=Bottleneck)	1
layers	list[int] | None	List of block counts per layer (default: [3, 4, 6, 3] for ResNet-50)	None
base_channels	int	Fixed channel width for all layers (default: 64)	64
kernel_size	int	Convolution kernel size (default: 5)	5
kernel_size_stem	int	Stem convolution kernel size (default: 5)	5
stem_channels	tuple[int, ...]	Tuple of stem layer channels	(32, 32, 64)
dropout_rate	float	Dropout probability in head (default: 0.5)	0.5
concat_pooling	bool	Use avg+max concat pooling (default: True)	True
block_channels	tuple[int, ...] | None	Per-layer channel widths (before expansion); overrides the fixed base_channels width when given	None

Example

model = XResNet1dBenchmark(input_channels=12, output_size=5) x = torch.randn(32, 12, 1000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 128) with concat_pooling

Source code in deepecgkit/models/xresnet1d_benchmark.py

@register_model(
    name="xresnet1d-benchmark",
    description="XResNet1d from PTB-XL benchmark (fixed feature dim, concat pooling)",
)
class XResNet1dBenchmark(ECGModelBase):
    """
    XResNet1d adapted from the PTB-XL benchmarking repository.

    Key differences from the standard XResNet1D:
    - Fixed feature dimension across all layers (all blocks use same width)
    - Concat pooling head (AdaptiveAvgPool + AdaptiveMaxPool concatenated)
    - Multi-layer stem with configurable kernel size
    - Supports both BasicBlock (expansion=1) and Bottleneck (expansion=4)

    Reference:
        Strodthoff N., Wagner P., Schaeffter T., Samek W.
        "Deep Learning for ECG Analysis: Benchmarks and Insights from PTB-XL" (2021)
        https://github.com/helme/ecg_ptbxl_benchmarking

    Args:
        input_channels: Number of input channels (default: 12 for 12-lead ECG)
        output_size: Number of output classes
        expansion: Block expansion factor (1=BasicBlock, 4=Bottleneck)
        layers: List of block counts per layer (default: [3, 4, 6, 3] for ResNet-50)
        base_channels: Fixed channel width for all layers (default: 64)
        kernel_size: Convolution kernel size (default: 5)
        kernel_size_stem: Stem convolution kernel size (default: 5)
        stem_channels: Tuple of stem layer channels
        dropout_rate: Dropout probability in head (default: 0.5)
        concat_pooling: Use avg+max concat pooling (default: True)
        block_channels: Per-layer channel widths (before expansion); overrides
            the fixed ``base_channels`` width when given

    Example:
        >>> model = XResNet1dBenchmark(input_channels=12, output_size=5)
        >>> x = torch.randn(32, 12, 1000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 128) with concat_pooling
    """

    HEAD_ATTR: ClassVar[str] = "head"

    def __init__(
        self,
        input_channels: int = 12,
        output_size: int = 5,
        expansion: int = 1,
        layers: list[int] | None = None,
        base_channels: int = 64,
        kernel_size: int = 5,
        kernel_size_stem: int = 5,
        stem_channels: tuple[int, ...] = (32, 32, 64),
        dropout_rate: float = 0.5,
        concat_pooling: bool = True,
        block_channels: tuple[int, ...] | None = None,
    ):
        super().__init__()

        if layers is None:
            layers = [2, 2, 2, 2]
        if block_channels is not None and len(block_channels) != len(layers):
            raise ValueError(
                f"block_channels must have one entry per layer, got {len(block_channels)} "
                f"for {len(layers)} layers"
            )

        self.expansion = expansion
        self.concat_pooling = concat_pooling

        stem_szs = [input_channels, *stem_channels]
        self.stem = nn.Sequential(
            *[
                ConvBnAct(
                    stem_szs[i], stem_szs[i + 1], ks=kernel_size_stem, stride=2 if i == 0 else 1
                )
                for i in range(len(stem_channels))
            ]
        )
        self.stem_pool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)

        if block_channels is None:
            block_channels = tuple([base_channels] * len(layers))
        block_szs = [64 // expansion, *block_channels]

        self.res_layers = nn.Sequential(
            *[
                self._make_layer(
                    ni=block_szs[i],
                    nf=block_szs[i + 1],
                    blocks=n_blocks,
                    stride=1 if i == 0 else 2,
                    kernel_size=kernel_size,
                )
                for i, n_blocks in enumerate(layers)
            ]
        )

        final_nf = block_szs[-1] * expansion
        self._feature_dim = final_nf * 2 if concat_pooling else final_nf

        if concat_pooling:
            self.pool = nn.ModuleList([nn.AdaptiveAvgPool1d(1), nn.AdaptiveMaxPool1d(1)])
        else:
            self.pool = nn.ModuleList([nn.AdaptiveAvgPool1d(1)])

        self.head = nn.Sequential(
            nn.Flatten(),
            nn.BatchNorm1d(self._feature_dim),
            nn.Dropout(dropout_rate / 2),
            nn.Linear(self._feature_dim, output_size),
        )

        _init_cnn(self)

    def _make_layer(
        self, ni: int, nf: int, blocks: int, stride: int, kernel_size: int
    ) -> nn.Sequential:
        return nn.Sequential(
            *[
                XResBlock(
                    self.expansion,
                    ni if i == 0 else nf,
                    nf,
                    stride=stride if i == 0 else 1,
                    kernel_size=kernel_size,
                )
                for i in range(blocks)
            ]
        )

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        x = self.stem_pool(x)
        x = self.res_layers(x)
        pooled = torch.cat([p(x) for p in self.pool], dim=1)
        return torch.flatten(pooled, 1)

KanResWideX

Bases: ECGModelBase

KanRes-Wide-X model for ECG signal classification.

A convolutional neural network architecture designed for ECG signal analysis with residual connections and wide blocks for improved feature extraction.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes or regression targets	4
base_channels	int	Base number of channels for the first layer (default: 64)	64

Example

model = KanResWideX(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape) # [32, 4]

features = model.extract_features(x) print(features.shape) # (32, 64)

Source code in deepecgkit/models/kanres_wide_x.py

@register_model(
    name="kanres",
    description="KAN-ResNet architecture with wide layers",
)
class KanResWideX(ECGModelBase):
    """
    KanRes-Wide-X model for ECG signal classification.

    A convolutional neural network architecture designed for ECG signal analysis
    with residual connections and wide blocks for improved feature extraction.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes or regression targets
        base_channels: Base number of channels for the first layer (default: 64)

    Example:
        >>> model = KanResWideX(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)  # [32, 4]

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 64)
    """

    def __init__(self, input_channels: int = 1, output_size: int = 4, base_channels: int = 64):
        super().__init__()

        self.input_layer = ConvBlock(input_channels, base_channels)
        self.res_modules = nn.Sequential(
            KanResModule(base_channels), KanResModule(base_channels), KanResModule(base_channels)
        )
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self._feature_dim = base_channels
        self.classifier = nn.Linear(base_channels, output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.input_layer(x)
        x = self.res_modules(x)
        x = self.global_pool(x)
        x = x.squeeze(-1)
        return x

KanResDeepX

Bases: ECGModelBase

KanRes-Deep-X model for ECG signal classification.

A deep residual convolutional neural network architecture designed for ECG signal analysis with 8 residual blocks for improved feature extraction.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes (default: 4)	4
base_channels	int	Base number of channels for the architecture (default: 32)	32

Example

model = KanResDeepX(input_channels=1, output_size=4, base_channels=32) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 32)

Source code in deepecgkit/models/kanres_deep_x.py

@register_model(
    name="kanres-deep",
    description="Deep KAN-ResNet architecture",
    min_input_length=11,
)
class KanResDeepX(ECGModelBase):
    """KanRes-Deep-X model for ECG signal classification.

    A deep residual convolutional neural network architecture designed for ECG
    signal analysis with 8 residual blocks for improved feature extraction.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes (default: 4)
        base_channels: Base number of channels for the architecture (default: 32)

    Example:
        >>> model = KanResDeepX(input_channels=1, output_size=4, base_channels=32)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 32)
    """

    HEAD_ATTR: ClassVar[str] = "fc"

    def __init__(self, input_channels: int = 1, output_size: int = 4, base_channels: int = 32):
        super().__init__()

        self.base_channels = base_channels
        init_channels = base_channels * 2

        self.init_block = KanResInit(input_channels, init_channels, base_channels, 8, 3, 1)
        self.pool = nn.AvgPool1d(kernel_size=2)

        self.res_modules = nn.ModuleList(
            [
                KanResModule(base_channels, base_channels * 2, base_channels, 50, 50, 1)
                for _ in range(8)
            ]
        )

        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self._feature_dim = base_channels
        self.fc = nn.Linear(base_channels, output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.init_block(x)
        x = self.pool(x)

        for res_module in self.res_modules:
            x = res_module(x)

        x = self.global_pool(x)
        x = x.squeeze(-1)
        return x

DeepResCNN

Bases: ECGModelBase

Deep Residual 2D CNN for ECG classification.

Faithful implementation of Elyamani et al. (2022). Uses Conv2d with (1, k) kernels and valid padding to process each ECG lead independently along the time axis, then fuses across leads with a (leads, 1) convolution at the end.

Input convention: (batch, leads, time) -- standard deepecg-kit format. Internally reshaped to (batch, 1, leads, time) for 2D convolution.

The classifier head includes L2 regularization matching the original Keras model. Call l2_regularization_loss() to obtain the penalty term.

Reference

https://github.com/HaneenElyamani/ECG-classification

Parameters:

Name	Type	Description	Default
input_channels	int	Number of ECG leads (default: 12)	12
output_size	int	Number of output classes (default: 5)	5
dropout_rate	float	Dropout probability in residual blocks (default: 0.1)	0.1

Example

model = DeepResCNN(input_channels=12, output_size=5) x = torch.randn(32, 12, 1000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 128)

Source code in deepecgkit/models/deep_res_cnn.py

@register_model(
    name="deep-res-cnn",
    description="Deep Residual 2D CNN (Elyamani et al. 2022) for multi-lead ECG classification",
    min_input_length=127,
)
class DeepResCNN(ECGModelBase):
    """
    Deep Residual 2D CNN for ECG classification.

    Faithful implementation of Elyamani et al. (2022). Uses Conv2d with (1, k)
    kernels and valid padding to process each ECG lead independently along the
    time axis, then fuses across leads with a (leads, 1) convolution at the end.

    Input convention: (batch, leads, time) -- standard deepecg-kit format.
    Internally reshaped to (batch, 1, leads, time) for 2D convolution.

    The classifier head includes L2 regularization matching the original Keras
    model. Call l2_regularization_loss() to obtain the penalty term.

    Reference:
        https://github.com/HaneenElyamani/ECG-classification

    Args:
        input_channels: Number of ECG leads (default: 12)
        output_size: Number of output classes (default: 5)
        dropout_rate: Dropout probability in residual blocks (default: 0.1)

    Example:
        >>> model = DeepResCNN(input_channels=12, output_size=5)
        >>> x = torch.randn(32, 12, 1000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 128)
    """

    HEAD_ATTR: ClassVar[str] = "fc_out"

    def __init__(
        self,
        input_channels: int = 12,
        output_size: int = 5,
        dropout_rate: float = 0.1,
    ):
        super().__init__()

        self.stem_conv = nn.Conv2d(1, 32, kernel_size=(1, 7))
        self.stem_bn = nn.BatchNorm2d(32)

        self.block1 = StemResBlock2D(32, 64, dropout_rate=dropout_rate)
        self.block2 = PreActResBlock2D(64, 64, dropout_rate=dropout_rate)
        self.block3 = PreActResBlock2D(64, 128, dropout_rate=dropout_rate)
        self.block4 = PreActResBlock2D(128, 128, dropout_rate=dropout_rate)

        self.input_channels = input_channels
        self.lead_fusion = nn.Conv2d(128, 128, kernel_size=(input_channels, 1))
        self.lead_fusion_bn = nn.BatchNorm2d(128)
        self.pool = nn.AdaptiveAvgPool2d(1)

        self._feature_dim = 128

        self.fc1 = nn.Linear(128, 128)
        self.fc1_bn = nn.BatchNorm1d(128)
        self.fc1_drop = nn.Dropout(dropout_rate)

        self.fc2 = nn.Linear(128, 64)
        self.fc2_bn = nn.BatchNorm1d(64)
        self.fc2_drop = nn.Dropout(0.15)

        self.fc_out = nn.Linear(64, output_size)

    def l2_regularization_loss(self) -> torch.Tensor:
        """Return the L2 penalty for the classifier head weights.

        Matches the Keras kernel_regularizer=L2(lambda) on the two Dense layers.
        Add this to the training loss for full equivalence with the original model.
        """
        return 0.005 * self.fc1.weight.pow(2).sum() + 0.009 * self.fc2.weight.pow(2).sum()

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        """Return the pooled representation preceding the classifier head.

        Args:
            x: Input of shape (batch, leads, time)

        Returns:
            Tensor of shape (batch, feature_dim).

        Raises:
            ValueError: If the lead count differs from the configured one. The
                lead-fusion convolution slides along the lead axis and the
                adaptive pool absorbs the remainder, so a mismatch would
                otherwise produce a correctly shaped but meaningless embedding.
        """
        if x.dim() != 3:
            raise ValueError(
                f"{type(self).__name__} expects input of shape (batch, leads, time), "
                f"got a {x.dim()}-D tensor of shape {tuple(x.shape)}."
            )
        if x.shape[1] != self.input_channels:
            raise ValueError(
                f"{type(self).__name__} was built for {self.input_channels} leads but "
                f"received {x.shape[1]}. Rebuild the model with "
                f"input_channels={x.shape[1]} or pass matching data."
            )

        x = x.unsqueeze(1)

        x = self.stem_conv(x)
        x = self.stem_bn(x)
        x = func.relu(x)

        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.block4(x)

        x = self.lead_fusion(x)
        x = self.lead_fusion_bn(x)
        x = func.relu(x)
        x = self.pool(x)
        return torch.flatten(x, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.extract_features(x)

        x = self.fc1(x)
        x = self.fc1_bn(x)
        x = func.relu(x)
        x = self.fc1_drop(x)

        x = self.fc2(x)
        x = self.fc2_bn(x)
        x = func.relu(x)
        x = self.fc2_drop(x)

        x = self.fc_out(x)
        return x

l2_regularization_loss

l2_regularization_loss() -> torch.Tensor

Return the L2 penalty for the classifier head weights.

Matches the Keras kernel_regularizer=L2(lambda) on the two Dense layers. Add this to the training loss for full equivalence with the original model.

Source code in deepecgkit/models/deep_res_cnn.py

def l2_regularization_loss(self) -> torch.Tensor:
    """Return the L2 penalty for the classifier head weights.

    Matches the Keras kernel_regularizer=L2(lambda) on the two Dense layers.
    Add this to the training loss for full equivalence with the original model.
    """
    return 0.005 * self.fc1.weight.pow(2).sum() + 0.009 * self.fc2.weight.pow(2).sum()

extract_features

extract_features(x: Tensor) -> torch.Tensor

Return the pooled representation preceding the classifier head.

Parameters:

Name	Type	Description	Default
x	Tensor	Input of shape (batch, leads, time)	required

Returns:

Type	Description
Tensor	Tensor of shape (batch, feature_dim).

Raises:

Type	Description
ValueError	If the lead count differs from the configured one. The lead-fusion convolution slides along the lead axis and the adaptive pool absorbs the remainder, so a mismatch would otherwise produce a correctly shaped but meaningless embedding.

Source code in deepecgkit/models/deep_res_cnn.py

def extract_features(self, x: torch.Tensor) -> torch.Tensor:
    """Return the pooled representation preceding the classifier head.

    Args:
        x: Input of shape (batch, leads, time)

    Returns:
        Tensor of shape (batch, feature_dim).

    Raises:
        ValueError: If the lead count differs from the configured one. The
            lead-fusion convolution slides along the lead axis and the
            adaptive pool absorbs the remainder, so a mismatch would
            otherwise produce a correctly shaped but meaningless embedding.
    """
    if x.dim() != 3:
        raise ValueError(
            f"{type(self).__name__} expects input of shape (batch, leads, time), "
            f"got a {x.dim()}-D tensor of shape {tuple(x.shape)}."
        )
    if x.shape[1] != self.input_channels:
        raise ValueError(
            f"{type(self).__name__} was built for {self.input_channels} leads but "
            f"received {x.shape[1]}. Rebuild the model with "
            f"input_channels={x.shape[1]} or pass matching data."
        )

    x = x.unsqueeze(1)

    x = self.stem_conv(x)
    x = self.stem_bn(x)
    x = func.relu(x)

    x = self.block1(x)
    x = self.block2(x)
    x = self.block3(x)
    x = self.block4(x)

    x = self.lead_fusion(x)
    x = self.lead_fusion_bn(x)
    x = func.relu(x)
    x = self.pool(x)
    return torch.flatten(x, 1)

Modern CNN

ConvNeXtV21D

Bases: ECGModelBase

ConvNeXtV2 adapted for 1D ECG signals.

A modern convolutional architecture using depthwise separable convolutions, LayerNorm, GELU activation, and Global Response Normalization (GRN).

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
dims	list | None	Channel dimensions for each stage (default: [64, 128, 256, 512])	None
depths	list | None	Number of blocks per stage (default: [2, 2, 6, 2])	None
kernel_size	int	Kernel size for depthwise convolutions (default: 7)	7
expansion_factor	int	Expansion ratio for inverted bottleneck (default: 4)	4
dropout_rate	float	Dropout probability (default: 0.3)	0.3

Example

model = ConvNeXtV21D(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 512)

Source code in deepecgkit/models/convnext_v2_1d.py

@register_model(
    name="convnext-v2",
    description="Modern ConvNet with depthwise convolutions and Global Response Normalization for ECG",
    min_input_length=29,
)
class ConvNeXtV21D(ECGModelBase):
    """
    ConvNeXtV2 adapted for 1D ECG signals.

    A modern convolutional architecture using depthwise separable convolutions,
    LayerNorm, GELU activation, and Global Response Normalization (GRN).

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        dims: Channel dimensions for each stage (default: [64, 128, 256, 512])
        depths: Number of blocks per stage (default: [2, 2, 6, 2])
        kernel_size: Kernel size for depthwise convolutions (default: 7)
        expansion_factor: Expansion ratio for inverted bottleneck (default: 4)
        dropout_rate: Dropout probability (default: 0.3)

    Example:
        >>> model = ConvNeXtV21D(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 512)
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        dims: list | None = None,
        depths: list | None = None,
        kernel_size: int = 7,
        expansion_factor: int = 4,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        if dims is None:
            dims = [64, 128, 256, 512]
        if depths is None:
            depths = [2, 2, 6, 2]

        self.stem = nn.Sequential(
            nn.Conv1d(input_channels, dims[0], kernel_size=7, stride=4, padding=3),
            LayerNorm1d(dims[0]),
        )

        stages = []
        for i in range(len(dims)):
            stage_blocks = []
            for _ in range(depths[i]):
                stage_blocks.append(ConvNeXtV2Block1D(dims[i], expansion_factor, kernel_size))
            stages.append(nn.Sequential(*stage_blocks))

        self.stages = nn.ModuleList(stages)

        self.downsamples = nn.ModuleList()
        for i in range(len(dims) - 1):
            self.downsamples.append(
                nn.Sequential(
                    LayerNorm1d(dims[i]),
                    nn.Conv1d(dims[i], dims[i + 1], kernel_size=2, stride=2),
                )
            )

        self.pool = nn.AdaptiveAvgPool1d(1)
        self.final_norm = nn.LayerNorm(dims[-1])
        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = dims[-1]

        self.classifier = nn.Linear(dims[-1], output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        for i, stage in enumerate(self.stages):
            x = stage(x)
            if i < len(self.downsamples):
                x = self.downsamples[i](x)
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.final_norm(x)
        return x

InceptionTime1D

Bases: ECGModelBase

InceptionTime model adapted for 1D ECG signal classification.

Uses parallel convolutions at multiple temporal scales (short/medium/long kernels) with residual connections to capture both rapid arrhythmia spikes and slow rhythm patterns simultaneously.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
n_filters	int	Number of filters per Inception branch (default: 32)	32
depth	int	Number of Inception residual blocks (default: 6)	6
kernel_sizes	tuple[int, ...]	Tuple of kernel sizes for multi-scale branches	(5, 15, 41)
bottleneck_channels	int	Channels in bottleneck layers (default: 32)	32
dropout_rate	float	Dropout probability (default: 0.3)	0.3

Example

model = InceptionTime1D(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 128) with n_filters=32

Source code in deepecgkit/models/inception_time.py

@register_model(
    name="inception-time",
    description="Multi-scale temporal CNN inspired by InceptionTime",
)
class InceptionTime1D(ECGModelBase):
    """
    InceptionTime model adapted for 1D ECG signal classification.

    Uses parallel convolutions at multiple temporal scales (short/medium/long
    kernels) with residual connections to capture both rapid arrhythmia spikes
    and slow rhythm patterns simultaneously.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        n_filters: Number of filters per Inception branch (default: 32)
        depth: Number of Inception residual blocks (default: 6)
        kernel_sizes: Tuple of kernel sizes for multi-scale branches
        bottleneck_channels: Channels in bottleneck layers (default: 32)
        dropout_rate: Dropout probability (default: 0.3)

    Example:
        >>> model = InceptionTime1D(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 128) with n_filters=32
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        n_filters: int = 32,
        depth: int = 6,
        kernel_sizes: tuple[int, ...] = (5, 15, 41),
        bottleneck_channels: int = 32,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        n_branches = len(kernel_sizes) + 1
        block_out_channels = n_filters * n_branches

        blocks = []
        for i in range(depth):
            in_ch = input_channels if i == 0 else block_out_channels
            blocks.append(
                InceptionResidualBlock1D(
                    in_channels=in_ch,
                    n_filters=n_filters,
                    kernel_sizes=kernel_sizes,
                    bottleneck_channels=bottleneck_channels,
                )
            )

        self.blocks = nn.Sequential(*blocks)
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = block_out_channels

        self.classifier = nn.Linear(block_out_channels, output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.blocks(x)
        x = self.global_pool(x)
        x = torch.flatten(x, 1)
        return x

TCN

Bases: ECGModelBase

Temporal Convolutional Network for ECG signal classification.

Uses stacked dilated causal convolutions with exponentially growing receptive fields to efficiently model long-range dependencies without recurrence. The causal structure makes it suitable for real-time ECG monitoring applications.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
num_channels	list | None	List of channel sizes per temporal block (default: [64, 64, 128, 128, 256, 256])	None
kernel_size	int	Convolution kernel size (default: 7)	7
dropout_rate	float	Dropout probability (default: 0.2)	0.2

Example

model = TCN(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 256) with default num_channels

Source code in deepecgkit/models/tcn.py

@register_model(
    name="tcn",
    description="Temporal Convolutional Network with dilated causal convolutions",
)
class TCN(ECGModelBase):
    """
    Temporal Convolutional Network for ECG signal classification.

    Uses stacked dilated causal convolutions with exponentially growing
    receptive fields to efficiently model long-range dependencies without
    recurrence. The causal structure makes it suitable for real-time
    ECG monitoring applications.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        num_channels: List of channel sizes per temporal block
            (default: [64, 64, 128, 128, 256, 256])
        kernel_size: Convolution kernel size (default: 7)
        dropout_rate: Dropout probability (default: 0.2)

    Example:
        >>> model = TCN(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 256) with default num_channels
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        num_channels: list | None = None,
        kernel_size: int = 7,
        dropout_rate: float = 0.2,
    ):
        super().__init__()

        if num_channels is None:
            num_channels = [64, 64, 128, 128, 256, 256]

        blocks = []
        for i, out_ch in enumerate(num_channels):
            in_ch = input_channels if i == 0 else num_channels[i - 1]
            dilation = 2**i
            blocks.append(
                TemporalBlock(
                    in_channels=in_ch,
                    out_channels=out_ch,
                    kernel_size=kernel_size,
                    dilation=dilation,
                    dropout_rate=dropout_rate,
                )
            )

        self.network = nn.Sequential(*blocks)
        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self._feature_dim = num_channels[-1]
        self.dropout = nn.Dropout(dropout_rate)
        self.classifier = nn.Linear(self._feature_dim, output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.network(x)
        x = self.global_pool(x)
        x = torch.flatten(x, 1)
        return x

Recurrent Networks

CRNN

Bases: ECGModelBase

Convolutional Recurrent Neural Network for ECG classification.

Uses a CNN front-end for local feature extraction followed by a bidirectional LSTM for temporal aggregation, combining the strengths of both architectures.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
cnn_channels	list | None	List of channel sizes for CNN stages (default: [32, 64, 128, 256])	None
lstm_hidden_size	int	Size of LSTM hidden state (default: 128)	128
lstm_num_layers	int	Number of LSTM layers (default: 2)	2
bidirectional	bool	Use bidirectional LSTM (default: True)	True
dropout_rate	float	Dropout probability (default: 0.3)	0.3

Example

model = CRNN(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 256)

Source code in deepecgkit/models/crnn.py

@register_model(
    name="crnn",
    description="CNN-LSTM hybrid for local feature extraction and temporal aggregation",
    min_input_length=16,
)
class CRNN(ECGModelBase):
    """
    Convolutional Recurrent Neural Network for ECG classification.

    Uses a CNN front-end for local feature extraction followed by a
    bidirectional LSTM for temporal aggregation, combining the strengths
    of both architectures.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        cnn_channels: List of channel sizes for CNN stages (default: [32, 64, 128, 256])
        lstm_hidden_size: Size of LSTM hidden state (default: 128)
        lstm_num_layers: Number of LSTM layers (default: 2)
        bidirectional: Use bidirectional LSTM (default: True)
        dropout_rate: Dropout probability (default: 0.3)

    Example:
        >>> model = CRNN(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 256)
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        cnn_channels: list | None = None,
        lstm_hidden_size: int = 128,
        lstm_num_layers: int = 2,
        bidirectional: bool = True,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        if cnn_channels is None:
            cnn_channels = [32, 64, 128, 256]

        self.lstm_hidden_size = lstm_hidden_size
        self.lstm_num_layers = lstm_num_layers
        self.bidirectional = bidirectional
        self.num_directions = 2 if bidirectional else 1

        layers = []
        in_ch = input_channels
        kernel_sizes = [7, 5, 3, 3]
        for i, out_ch in enumerate(cnn_channels):
            ks = kernel_sizes[i] if i < len(kernel_sizes) else 3
            layers.extend(
                [
                    nn.Conv1d(in_ch, out_ch, kernel_size=ks, padding=ks // 2),
                    nn.BatchNorm1d(out_ch),
                    nn.ReLU(),
                    nn.MaxPool1d(kernel_size=2),
                    nn.Dropout(dropout_rate),
                ]
            )
            in_ch = out_ch

        self.cnn = nn.Sequential(*layers)

        self.lstm = nn.LSTM(
            input_size=cnn_channels[-1],
            hidden_size=lstm_hidden_size,
            num_layers=lstm_num_layers,
            batch_first=True,
            dropout=dropout_rate if lstm_num_layers > 1 else 0,
            bidirectional=bidirectional,
        )

        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = lstm_hidden_size * self.num_directions

        self.classifier = nn.Sequential(
            nn.Linear(self._feature_dim, 64),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(64, output_size),
        )

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        batch_size = x.size(0)
        x = self.cnn(x)
        x = x.transpose(1, 2)
        _, (h_n, _) = self.lstm(x)

        return lstm_final_hidden(
            h_n,
            num_layers=self.lstm_num_layers,
            batch_size=batch_size,
            hidden_size=self.lstm_hidden_size,
            bidirectional=self.bidirectional,
        )

GRUECG

Bases: ECGModelBase

GRU-based model for ECG signal classification.

A recurrent neural network using GRU layers with concat pooling (adaptive avg + adaptive max + last hidden state) for ECG classification. Based on the RNN architecture from the PTB-XL benchmark.

Reference

https://github.com/helme/ecg_ptbxl_benchmarking

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 12 for 12-lead ECG)	12
output_size	int	Number of output classes	5
hidden_size	int	Size of GRU hidden state (default: 256)	256
num_layers	int	Number of GRU layers (default: 2)	2
dropout_rate	float	Dropout probability (default: 0.3)	0.3
bidirectional	bool	Use bidirectional GRU (default: False)	False

Example

model = GRUECG(input_channels=12, output_size=5) x = torch.randn(32, 12, 1000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 768) with hidden_size=256, unidirectional

Source code in deepecgkit/models/gru_model.py

@register_model(
    name="gru",
    description="GRU-based model for sequential ECG analysis",
)
class GRUECG(ECGModelBase):
    """
    GRU-based model for ECG signal classification.

    A recurrent neural network using GRU layers with concat pooling
    (adaptive avg + adaptive max + last hidden state) for ECG classification.
    Based on the RNN architecture from the PTB-XL benchmark.

    Reference:
        https://github.com/helme/ecg_ptbxl_benchmarking

    Args:
        input_channels: Number of input channels (default: 12 for 12-lead ECG)
        output_size: Number of output classes
        hidden_size: Size of GRU hidden state (default: 256)
        num_layers: Number of GRU layers (default: 2)
        dropout_rate: Dropout probability (default: 0.3)
        bidirectional: Use bidirectional GRU (default: False)

    Example:
        >>> model = GRUECG(input_channels=12, output_size=5)
        >>> x = torch.randn(32, 12, 1000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 768) with hidden_size=256, unidirectional
    """

    HEAD_ATTR: ClassVar[str] = "head"

    def __init__(
        self,
        input_channels: int = 12,
        output_size: int = 5,
        hidden_size: int = 256,
        num_layers: int = 2,
        dropout_rate: float = 0.3,
        bidirectional: bool = False,
    ):
        super().__init__()

        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bidirectional = bidirectional
        self.num_directions = 2 if bidirectional else 1

        self.gru = nn.GRU(
            input_size=input_channels,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout_rate if num_layers > 1 else 0,
            bidirectional=bidirectional,
        )

        rnn_out_dim = hidden_size * self.num_directions
        self._feature_dim = rnn_out_dim * 3

        self.head = nn.Sequential(
            nn.BatchNorm1d(self._feature_dim),
            nn.Dropout(dropout_rate / 2),
            nn.Linear(self._feature_dim, output_size),
        )

    def _concat_pool(self, rnn_output: torch.Tensor) -> torch.Tensor:
        x = rnn_output.transpose(1, 2)
        avg_pool = torch.mean(x, dim=2)
        max_pool, _ = torch.max(x, dim=2)

        if self.bidirectional:
            last = torch.cat(
                [rnn_output[:, -1, : self.hidden_size], rnn_output[:, 0, self.hidden_size :]], dim=1
            )
        else:
            last = rnn_output[:, -1, :]

        return torch.cat([avg_pool, max_pool, last], dim=1)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = x.transpose(1, 2)
        output, _ = self.gru(x)
        return self._concat_pool(output)

LSTMECG

Bases: ECGModelBase

LSTM-based model for ECG signal classification.

A recurrent neural network using bidirectional LSTM layers for temporal pattern recognition in ECG signals.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
hidden_size	int	Size of LSTM hidden state (default: 128)	128
num_layers	int	Number of LSTM layers (default: 2)	2
dropout_rate	float	Dropout probability (default: 0.3)	0.3
bidirectional	bool	Use bidirectional LSTM (default: True)	True

Example

model = LSTMECG(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 256) with bidirectional=True

Source code in deepecgkit/models/lstm_model.py

@register_model(
    name="lstm",
    description="LSTM-based model for sequential ECG analysis",
)
class LSTMECG(ECGModelBase):
    """
    LSTM-based model for ECG signal classification.

    A recurrent neural network using bidirectional LSTM layers
    for temporal pattern recognition in ECG signals.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        hidden_size: Size of LSTM hidden state (default: 128)
        num_layers: Number of LSTM layers (default: 2)
        dropout_rate: Dropout probability (default: 0.3)
        bidirectional: Use bidirectional LSTM (default: True)

    Example:
        >>> model = LSTMECG(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 256) with bidirectional=True
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        hidden_size: int = 128,
        num_layers: int = 2,
        dropout_rate: float = 0.3,
        bidirectional: bool = True,
    ):
        super().__init__()

        self.input_channels = input_channels
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.bidirectional = bidirectional
        self.num_directions = 2 if bidirectional else 1

        self.lstm = nn.LSTM(
            input_size=input_channels,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout_rate if num_layers > 1 else 0,
            bidirectional=bidirectional,
        )

        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = hidden_size * self.num_directions

        self.classifier = nn.Sequential(
            nn.Linear(self._feature_dim, 64),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(64, output_size),
        )

    def _extract_lstm_features(self, x: torch.Tensor) -> torch.Tensor:
        batch_size = x.size(0)
        x = x.transpose(1, 2)
        _, (h_n, _) = self.lstm(x)

        return lstm_final_hidden(
            h_n,
            num_layers=self.num_layers,
            batch_size=batch_size,
            hidden_size=self.hidden_size,
            bidirectional=self.bidirectional,
        )

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        return self._extract_lstm_features(x)

Transformers & Attention

TransformerECG

Bases: ECGModelBase

Transformer-based model for ECG signal classification.

A transformer architecture that uses self-attention mechanisms to capture long-range dependencies in ECG signals.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
d_model	int	Dimension of the model (default: 128)	128
nhead	int	Number of attention heads (default: 8)	8
num_encoder_layers	int	Number of transformer encoder layers (default: 4)	4
dim_feedforward	int	Dimension of feedforward network (default: 512)	512
dropout_rate	float	Dropout probability (default: 0.1)	0.1
max_len	int	Maximum sequence length (default: 5000)	5000

Example

model = TransformerECG(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 128) with d_model=128

Source code in deepecgkit/models/transformer_ecg.py

@register_model(
    name="transformer",
    description="Transformer-based ECG classifier",
    max_input_length=5000,
)
class TransformerECG(ECGModelBase):
    """
    Transformer-based model for ECG signal classification.

    A transformer architecture that uses self-attention mechanisms
    to capture long-range dependencies in ECG signals.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        d_model: Dimension of the model (default: 128)
        nhead: Number of attention heads (default: 8)
        num_encoder_layers: Number of transformer encoder layers (default: 4)
        dim_feedforward: Dimension of feedforward network (default: 512)
        dropout_rate: Dropout probability (default: 0.1)
        max_len: Maximum sequence length (default: 5000)

    Example:
        >>> model = TransformerECG(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 128) with d_model=128
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        d_model: int = 128,
        nhead: int = 8,
        num_encoder_layers: int = 4,
        dim_feedforward: int = 512,
        dropout_rate: float = 0.1,
        max_len: int = 5000,
    ):
        super().__init__()

        self.d_model = d_model

        self.input_projection = nn.Linear(input_channels, d_model)

        self.positional_encoding = PositionalEncoding(d_model, max_len, dropout_rate)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout_rate,
            batch_first=True,
        )

        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer,
            num_layers=num_encoder_layers,
        )

        self.global_pool = nn.AdaptiveAvgPool1d(1)
        self._feature_dim = d_model

        self.classifier = nn.Sequential(
            nn.Linear(d_model, dim_feedforward // 2),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(dim_feedforward // 2, output_size),
        )

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = x.transpose(1, 2)
        x = self.input_projection(x)
        x = self.positional_encoding(x)
        x = self.transformer_encoder(x)
        x = x.transpose(1, 2)
        x = self.global_pool(x)
        x = x.squeeze(-1)
        return x

Medformer

Bases: ECGModelBase

Medformer: Multi-Granularity Patching Transformer for medical time series.

Uses multiple patch sizes to capture fine, medium, and coarse temporal patterns, with intra-granularity self-attention and inter-granularity cross-attention for information fusion.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
d_model	int	Transformer model dimension (default: 128)	128
patch_sizes	tuple[int, ...]	Tuple of patch sizes for different granularities (default: (10, 25, 50))	(10, 25, 50)
num_encoder_layers	int	Number of encoder layers (default: 2)	2
nhead	int	Number of attention heads (default: 8)	8
dim_feedforward	int	Feedforward dimension (default: 256)	256
dropout_rate	float	Dropout probability (default: 0.1)	0.1
max_patches	int	Maximum number of patches per granularity (default: 500)	500

Example

model = Medformer(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 384)

Source code in deepecgkit/models/medformer.py

@register_model(
    name="medformer",
    description="Multi-granularity patching Transformer for medical time series (NeurIPS 2024)",
    min_input_length=50,
    max_input_length=5009,
)
class Medformer(ECGModelBase):
    """
    Medformer: Multi-Granularity Patching Transformer for medical time series.

    Uses multiple patch sizes to capture fine, medium, and coarse temporal
    patterns, with intra-granularity self-attention and inter-granularity
    cross-attention for information fusion.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        d_model: Transformer model dimension (default: 128)
        patch_sizes: Tuple of patch sizes for different granularities (default: (10, 25, 50))
        num_encoder_layers: Number of encoder layers (default: 2)
        nhead: Number of attention heads (default: 8)
        dim_feedforward: Feedforward dimension (default: 256)
        dropout_rate: Dropout probability (default: 0.1)
        max_patches: Maximum number of patches per granularity (default: 500)

    Example:
        >>> model = Medformer(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 384)
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        d_model: int = 128,
        patch_sizes: tuple[int, ...] = (10, 25, 50),
        num_encoder_layers: int = 2,
        nhead: int = 8,
        dim_feedforward: int = 256,
        dropout_rate: float = 0.1,
        max_patches: int = 500,
    ):
        super().__init__()

        self.patch_sizes = patch_sizes
        self.d_model = d_model
        num_granularities = len(patch_sizes)

        self.patch_embeddings = nn.ModuleList(
            [PatchEmbedding1D(input_channels, d_model, ps) for ps in patch_sizes]
        )

        self.pos_embeddings = nn.ParameterList(
            [nn.Parameter(torch.randn(1, max_patches, d_model) * 0.02) for _ in patch_sizes]
        )

        self.cross_channel_attn = CrossChannelAttention(
            d_model,
            nhead=min(4, nhead),
            dropout=dropout_rate,
        )

        self.intra_blocks = nn.ModuleList()
        self.inter_blocks = nn.ModuleList()
        for _ in range(num_encoder_layers):
            self.intra_blocks.append(
                nn.ModuleList(
                    [
                        IntraGranularityBlock(d_model, nhead, dim_feedforward, dropout_rate)
                        for _ in patch_sizes
                    ]
                )
            )
            self.inter_blocks.append(
                InterGranularityAttention(d_model, num_granularities, min(4, nhead), dropout_rate)
            )

        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = d_model * num_granularities

        self.classifier = nn.Linear(self._feature_dim, output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        granularity_features = []
        for patch_emb, pos_emb in zip(self.patch_embeddings, self.pos_embeddings, strict=False):
            patches = patch_emb(x)
            num_patches = patches.size(1)
            if num_patches > pos_emb.size(1):
                raise ValueError(
                    f"Medformer received {x.shape[-1]} samples per lead, which makes "
                    f"{num_patches} patches of {patch_emb.proj.kernel_size[0]} samples, but "
                    f"was built for {pos_emb.size(1)} per patch size. Pass shorter signals (a "
                    f"shorter segment or a lower sampling rate), or build the model with a "
                    f"larger max_patches."
                )
            patches = patches + pos_emb[:, :num_patches]
            granularity_features.append(patches)

        for intra_layer, inter_layer in zip(self.intra_blocks, self.inter_blocks, strict=False):
            for i, intra_block in enumerate(intra_layer):
                granularity_features[i] = intra_block(granularity_features[i])
            granularity_features = inter_layer(granularity_features)

        pooled = [f.mean(dim=1) for f in granularity_features]
        return torch.cat(pooled, dim=1)

Hybrid & State Space

ECGDualNet

Bases: ECGModelBase

Dual-path ECG classification network.

Runs a CNN-LSTM branch and a Transformer branch in parallel, then fuses their outputs via concatenation and a fusion head.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
cnn_channels	int	CNN output channels in CNN-LSTM branch (default: 128)	128
lstm_hidden	int	LSTM hidden size in CNN-LSTM branch (default: 128)	128
lstm_layers	int	Number of LSTM layers (default: 1)	1
d_model	int	Transformer model dimension (default: 128)	128
nhead	int	Number of attention heads (default: 8)	8
transformer_layers	int	Number of transformer encoder layers (default: 2)	2
dim_feedforward	int	Transformer feedforward dimension (default: 256)	256
dropout_rate	float	Dropout probability (default: 0.3)	0.3

Example

model = ECGDualNet(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 384)

Source code in deepecgkit/models/ecg_dualnet.py

@register_model(
    name="dualnet",
    description="Dual-path architecture with CNN-LSTM and Transformer branches for ECG",
    min_input_length=8,
    max_input_length=5000,
)
class ECGDualNet(ECGModelBase):
    """
    Dual-path ECG classification network.

    Runs a CNN-LSTM branch and a Transformer branch in parallel,
    then fuses their outputs via concatenation and a fusion head.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        cnn_channels: CNN output channels in CNN-LSTM branch (default: 128)
        lstm_hidden: LSTM hidden size in CNN-LSTM branch (default: 128)
        lstm_layers: Number of LSTM layers (default: 1)
        d_model: Transformer model dimension (default: 128)
        nhead: Number of attention heads (default: 8)
        transformer_layers: Number of transformer encoder layers (default: 2)
        dim_feedforward: Transformer feedforward dimension (default: 256)
        dropout_rate: Dropout probability (default: 0.3)

    Example:
        >>> model = ECGDualNet(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 384)
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        cnn_channels: int = 128,
        lstm_hidden: int = 128,
        lstm_layers: int = 1,
        d_model: int = 128,
        nhead: int = 8,
        transformer_layers: int = 2,
        dim_feedforward: int = 256,
        dropout_rate: float = 0.3,
    ):
        super().__init__()

        self.cnn_lstm_branch = CNNLSTMBranch(
            input_channels=input_channels,
            cnn_channels=cnn_channels,
            lstm_hidden=lstm_hidden,
            lstm_layers=lstm_layers,
            dropout_rate=dropout_rate,
        )

        self.transformer_branch = TransformerBranch(
            input_channels=input_channels,
            d_model=d_model,
            nhead=nhead,
            num_layers=transformer_layers,
            dim_feedforward=dim_feedforward,
            dropout_rate=dropout_rate,
        )

        self._feature_dim = self.cnn_lstm_branch.output_dim + self.transformer_branch.output_dim

        self.classifier = nn.Sequential(
            nn.Linear(self._feature_dim, 128),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(128, output_size),
        )

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        cnn_lstm_out = self.cnn_lstm_branch(x)
        transformer_out = self.transformer_branch(x)
        return torch.cat([cnn_lstm_out, transformer_out], dim=1)

Mamba1D

Bases: ECGModelBase

Bidirectional Mamba model for ECG classification.

Uses selective state space models with linear complexity as an alternative to Transformer self-attention. The SSM is implemented from scratch in pure PyTorch with no external dependencies.

Parameters:

Name	Type	Description	Default
input_channels	int	Number of input channels (default: 1 for single-lead ECG)	1
output_size	int	Number of output classes	4
d_model	int	Model dimension (default: 128)	128
d_state	int	State space dimension (default: 16)	16
d_conv	int	Local convolution width (default: 4)	4
expansion_factor	int	Inner dimension expansion (default: 2)	2
num_layers	int	Number of bidirectional Mamba layers (default: 4)	4
patch_size	int	Patch size for input tokenization (default: 50)	50
dropout_rate	float	Dropout probability (default: 0.1)	0.1
max_patches	int	Maximum number of patches (default: 500)	500

Example

model = Mamba1D(input_channels=1, output_size=4) x = torch.randn(32, 1, 3000) output = model(x) print(output.shape)

features = model.extract_features(x) print(features.shape) # (32, 128)

Source code in deepecgkit/models/mamba1d.py

@register_model(
    name="mamba",
    description="Bidirectional Mamba state space model with linear complexity for ECG",
    min_input_length=50,
    max_input_length=25049,
)
class Mamba1D(ECGModelBase):
    """
    Bidirectional Mamba model for ECG classification.

    Uses selective state space models with linear complexity as an
    alternative to Transformer self-attention. The SSM is implemented
    from scratch in pure PyTorch with no external dependencies.

    Args:
        input_channels: Number of input channels (default: 1 for single-lead ECG)
        output_size: Number of output classes
        d_model: Model dimension (default: 128)
        d_state: State space dimension (default: 16)
        d_conv: Local convolution width (default: 4)
        expansion_factor: Inner dimension expansion (default: 2)
        num_layers: Number of bidirectional Mamba layers (default: 4)
        patch_size: Patch size for input tokenization (default: 50)
        dropout_rate: Dropout probability (default: 0.1)
        max_patches: Maximum number of patches (default: 500)

    Example:
        >>> model = Mamba1D(input_channels=1, output_size=4)
        >>> x = torch.randn(32, 1, 3000)
        >>> output = model(x)
        >>> print(output.shape)

        >>> features = model.extract_features(x)
        >>> print(features.shape)  # (32, 128)
    """

    def __init__(
        self,
        input_channels: int = 1,
        output_size: int = 4,
        d_model: int = 128,
        d_state: int = 16,
        d_conv: int = 4,
        expansion_factor: int = 2,
        num_layers: int = 4,
        patch_size: int = 50,
        dropout_rate: float = 0.1,
        max_patches: int = 500,
    ):
        super().__init__()

        self.patch_embed = nn.Conv1d(
            input_channels,
            d_model,
            kernel_size=patch_size,
            stride=patch_size,
        )
        self.pos_embedding = nn.Parameter(torch.randn(1, max_patches, d_model) * 0.02)

        self.layers = nn.ModuleList(
            [
                BidirectionalMamba(d_model, d_state, d_conv, expansion_factor)
                for _ in range(num_layers)
            ]
        )

        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout_rate)
        self._feature_dim = d_model

        self.classifier = nn.Linear(d_model, output_size)

    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.patch_embed(x)
        x = x.transpose(1, 2)
        num_patches = x.size(1)
        x = x + self.pos_embedding[:, :num_patches]

        for layer in self.layers:
            x = layer(x)

        x = self.norm(x)
        x = x.mean(dim=1)
        return x

Pretrained Weights

PretrainedMixin

Adds pretrained weight loading to a registered model.

Mix into any model whose constructor follows the shared (input_channels, output_size, **kwargs) convention.

Example

model = KanResWideX.from_pretrained("kanres-af-30s") model = KanResWideX.from_pretrained("kanres-af-30s", map_location="cuda") model = KanResWideX.from_pretrained("/path/to/weights.pt", output_size=2) model = AFModel.from_pretrained("/path/to/trained.ckpt")

Source code in deepecgkit/models/pretrained.py

class PretrainedMixin:
    """Adds pretrained weight loading to a registered model.

    Mix into any model whose constructor follows the shared
    ``(input_channels, output_size, **kwargs)`` convention.

    Example:
        >>> model = KanResWideX.from_pretrained("kanres-af-30s")
        >>> model = KanResWideX.from_pretrained("kanres-af-30s", map_location="cuda")
        >>> model = KanResWideX.from_pretrained("/path/to/weights.pt", output_size=2)
        >>> model = AFModel.from_pretrained("/path/to/trained.ckpt")
    """

    @classmethod
    def from_pretrained(
        cls: type[T],
        weights: str | Path,
        map_location: str | torch.device | None = None,
        force_download: bool = False,
        **kwargs: Any,
    ) -> T:
        """Build the model and load pretrained weights into it.

        Args:
            weights: Name of registered pretrained weights, or a path to a
                weights file, an ECGTrainer checkpoint, or a legacy PyTorch
                Lightning checkpoint
            map_location: Device to map the weights onto while loading
            force_download: Re-download registered weights even if cached
            **kwargs: Constructor arguments, overriding anything recorded in the
                checkpoint or the weight registry

        Returns:
            The model with the weights loaded.

        Raises:
            ValueError: If the weights name is unknown, or the file holds no
                recognisable weights.
        """
        weight_path = Path(weights)
        if weight_path.exists():
            loaded = torch.load(weight_path, map_location=map_location, weights_only=True)
            state_dict, recorded = _split_checkpoint(loaded)
            model = cls(**{**recorded, **kwargs})
        else:
            info = get_weight_info(str(weights))
            model = cls(**{**info["model_kwargs"], **kwargs})
            loaded = load_pretrained_weights(str(weights), map_location, force_download)
            state_dict, _ = _split_checkpoint(loaded)

        model.load_state_dict(state_dict)
        return model

from_pretrained classmethod

from_pretrained(
    weights: str | Path,
    map_location: str | device | None = None,
    force_download: bool = False,
    **kwargs: Any,
) -> T

Build the model and load pretrained weights into it.

Parameters:

Name	Type	Description	Default
weights	str | Path	Name of registered pretrained weights, or a path to a weights file, an ECGTrainer checkpoint, or a legacy PyTorch Lightning checkpoint	required
map_location	str | device | None	Device to map the weights onto while loading	None
force_download	bool	Re-download registered weights even if cached	False
**kwargs	Any	Constructor arguments, overriding anything recorded in the checkpoint or the weight registry	{}

Returns:

Type	Description
T	The model with the weights loaded.

Raises:

Type	Description
ValueError	If the weights name is unknown, or the file holds no recognisable weights.

Source code in deepecgkit/models/pretrained.py

@classmethod
def from_pretrained(
    cls: type[T],
    weights: str | Path,
    map_location: str | torch.device | None = None,
    force_download: bool = False,
    **kwargs: Any,
) -> T:
    """Build the model and load pretrained weights into it.

    Args:
        weights: Name of registered pretrained weights, or a path to a
            weights file, an ECGTrainer checkpoint, or a legacy PyTorch
            Lightning checkpoint
        map_location: Device to map the weights onto while loading
        force_download: Re-download registered weights even if cached
        **kwargs: Constructor arguments, overriding anything recorded in the
            checkpoint or the weight registry

    Returns:
        The model with the weights loaded.

    Raises:
        ValueError: If the weights name is unknown, or the file holds no
            recognisable weights.
    """
    weight_path = Path(weights)
    if weight_path.exists():
        loaded = torch.load(weight_path, map_location=map_location, weights_only=True)
        state_dict, recorded = _split_checkpoint(loaded)
        model = cls(**{**recorded, **kwargs})
    else:
        info = get_weight_info(str(weights))
        model = cls(**{**info["model_kwargs"], **kwargs})
        loaded = load_pretrained_weights(str(weights), map_location, force_download)
        state_dict, _ = _split_checkpoint(loaded)

    model.load_state_dict(state_dict)
    return model

from_pretrained

from_pretrained(
    model_name: str,
    weights: str | Path,
    map_location: str | device | None = None,
    force_download: bool = False,
    **kwargs: Any,
) -> nn.Module

Build a registered model by name and load pretrained weights into it.

The callers that hold a model name rather than a class, such as the CLI and the feature extractor, use this instead of the classmethod.

Parameters:

Name	Type	Description	Default
model_name	str	Registered model name	required
weights	str | Path	Name of registered pretrained weights, or a path to a weights file or checkpoint	required
map_location	str | device | None	Device to map the weights onto while loading	None
force_download	bool	Re-download registered weights even if cached	False
**kwargs	Any	Constructor arguments forwarded to the model	{}

Returns:

Type	Description
Module	The model with the weights loaded.

Raises:

Type	Description
KeyError	If the model name is not registered.
TypeError	If the model does not support pretrained weight loading.

Source code in deepecgkit/models/pretrained.py

def from_pretrained(
    model_name: str,
    weights: str | Path,
    map_location: str | torch.device | None = None,
    force_download: bool = False,
    **kwargs: Any,
) -> nn.Module:
    """Build a registered model by name and load pretrained weights into it.

    The callers that hold a model name rather than a class, such as the CLI and
    the feature extractor, use this instead of the classmethod.

    Args:
        model_name: Registered model name
        weights: Name of registered pretrained weights, or a path to a weights
            file or checkpoint
        map_location: Device to map the weights onto while loading
        force_download: Re-download registered weights even if cached
        **kwargs: Constructor arguments forwarded to the model

    Returns:
        The model with the weights loaded.

    Raises:
        KeyError: If the model name is not registered.
        TypeError: If the model does not support pretrained weight loading.
    """
    model_class = get_model(model_name)
    if not hasattr(model_class, "from_pretrained"):
        raise TypeError(f"Model '{model_name}' does not support loading pretrained weights.")
    return model_class.from_pretrained(
        weights, map_location=map_location, force_download=force_download, **kwargs
    )

model_name_for_weights

model_name_for_weights(weight_name: str) -> str | None

Return the registry name of the model a set of pretrained weights targets.

The weight registry records the model's class name while the model registry is keyed on its short registry name, so the two need bridging.

Parameters:

Name	Type	Description	Default
weight_name	str	Key into the pretrained weight registry	required

Returns:

Type	Description
str | None	The registered model name, or None if the weights are unknown or their
str | None	class is not registered.

Source code in deepecgkit/models/pretrained.py

def model_name_for_weights(weight_name: str) -> str | None:
    """Return the registry name of the model a set of pretrained weights targets.

    The weight registry records the model's class name while the model registry
    is keyed on its short registry name, so the two need bridging.

    Args:
        weight_name: Key into the pretrained weight registry

    Returns:
        The registered model name, or None if the weights are unknown or their
        class is not registered.
    """
    entry = WEIGHTS_REGISTRY.get(weight_name)
    if entry is None:
        return None
    class_name = entry.get("model_class")
    for name in get_model_names():
        if get_model(name).__name__ == class_name:
            return name
    return None