deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/models/transformer_ecg.py

@@ -0,0 +1,127 @@
+import math
+
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class PositionalEncoding(nn.Module):
+    """
+    Positional encoding for transformer model.
+
+    Args:
+        d_model: Dimension of the model
+        max_len: Maximum sequence length (default: 5000)
+        dropout: Dropout probability (default: 0.1)
+    """
+
+    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(1, max_len, d_model)
+        pe[0, :, 0::2] = torch.sin(position * div_term)
+        pe[0, :, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        x = x + self.pe[:, : x.size(1)]
+        return self.dropout(x)
+
+
+@register_model(
+    name="transformer",
+    description="Transformer-based ECG classifier",
+)
+class TransformerECG(nn.Module):
+    """
+    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),
+        )
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def _extract_encoder_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
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        return self._extract_encoder_features(x)
+
+    def forward(self, x):
+        x = self.extract_features(x)
+        return self.classifier(x)

deepecgkit/models/xresnet1d.py

@@ -0,0 +1,265 @@
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class Mish(nn.Module):
+    """Mish activation: x * tanh(softplus(x))."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.tanh(func.softplus(x))
+
+
+class BlurPool1D(nn.Module):
+    """
+    Anti-aliased downsampling via blur-then-subsample.
+
+    Applies a fixed triangular low-pass filter before strided subsampling
+    to reduce aliasing artifacts during spatial reduction.
+
+    Args:
+        channels: Number of channels (applied per-channel)
+        stride: Downsampling stride (default: 2)
+    """
+
+    def __init__(self, channels: int, stride: int = 2):
+        super().__init__()
+        self.stride = stride
+        kernel = torch.tensor([1.0, 2.0, 1.0]) / 4.0
+        kernel = kernel.unsqueeze(0).unsqueeze(0).repeat(channels, 1, 1)
+        self.register_buffer("kernel", kernel)
+        self.channels = channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = func.pad(x, (1, 1), mode="reflect")
+        x = func.conv1d(x, self.kernel, stride=self.stride, groups=self.channels)
+        return x
+
+
+class XResBlock1D(nn.Module):
+    """
+    Improved residual block with Mish activation and optional blur-pool downsampling.
+
+    Args:
+        in_channels: Number of input channels
+        out_channels: Number of output channels
+        kernel_size: Convolution kernel size (default: 7)
+        stride: Convolution stride (default: 1)
+        downsample: Downsample layer for skip connection if needed
+        use_blur_pool: Whether to use anti-aliased downsampling
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 7,
+        stride: int = 1,
+        downsample: nn.Module | None = None,
+        use_blur_pool: bool = True,
+    ):
+        super().__init__()
+
+        self.conv1 = nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=kernel_size // 2,
+            bias=False,
+        )
+        self.bn1 = nn.BatchNorm1d(out_channels)
+
+        self.conv2 = nn.Conv1d(
+            out_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=kernel_size // 2,
+            bias=False,
+        )
+        self.bn2 = nn.BatchNorm1d(out_channels)
+
+        self.act = Mish()
+        self.downsample = downsample
+
+        if stride > 1 and use_blur_pool:
+            self.pool = BlurPool1D(out_channels, stride=stride)
+        elif stride > 1:
+            self.pool = nn.AvgPool1d(kernel_size=stride, stride=stride)
+        else:
+            self.pool = nn.Identity()
+
+        nn.init.zeros_(self.bn2.weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = self.act(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out = self.pool(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.act(out)
+        return out
+
+
+@register_model(
+    name="xresnet",
+    description="Improved ResNet with Mish activation and blur-pool downsampling",
+)
+class XResNet1D(nn.Module):
+    """
+    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
+    """
+
+    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):
+        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)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    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
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        x = self.fc(x)
+        return x

deepecgkit/models/xresnet1d_benchmark.py

@@ -0,0 +1,227 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+def _init_cnn(m: nn.Module):
+    if getattr(m, "bias", None) is not None:
+        nn.init.constant_(m.bias, 0)
+    if isinstance(m, (nn.Conv1d, nn.Linear)):
+        nn.init.kaiming_normal_(m.weight)
+    for child in m.children():
+        _init_cnn(child)
+
+
+class ConvBnAct(nn.Sequential):
+    """Conv1d → BatchNorm1d → Activation."""
+
+    def __init__(
+        self,
+        ni: int,
+        nf: int,
+        ks: int = 3,
+        stride: int = 1,
+        act: bool = True,
+        zero_bn: bool = False,
+    ):
+        layers: list[nn.Module] = [
+            nn.Conv1d(ni, nf, kernel_size=ks, stride=stride, padding=(ks - 1) // 2, bias=False),
+        ]
+        bn = nn.BatchNorm1d(nf)
+        if zero_bn:
+            nn.init.zeros_(bn.weight)
+        layers.append(bn)
+        if act:
+            layers.append(nn.ReLU(inplace=True))
+        super().__init__(*layers)
+
+
+class XResBlock(nn.Module):
+    """
+    XResNet residual block with expansion support.
+
+    Args:
+        expansion: Channel expansion factor (1 for BasicBlock, 4 for Bottleneck)
+        ni: Input channels (before expansion)
+        nf: Output channels (before expansion)
+        stride: Stride for downsampling
+        kernel_size: Convolution kernel size
+    """
+
+    def __init__(self, expansion: int, ni: int, nf: int, stride: int = 1, kernel_size: int = 5):
+        super().__init__()
+        nf_out = nf * expansion
+        ni_in = ni * expansion
+
+        if expansion == 1:
+            self.convs = nn.Sequential(
+                ConvBnAct(ni_in, nf, ks=kernel_size, stride=stride),
+                ConvBnAct(nf, nf_out, ks=kernel_size, zero_bn=True, act=False),
+            )
+        else:
+            self.convs = nn.Sequential(
+                ConvBnAct(ni_in, nf, ks=1),
+                ConvBnAct(nf, nf, ks=kernel_size, stride=stride),
+                ConvBnAct(nf, nf_out, ks=1, zero_bn=True, act=False),
+            )
+
+        id_layers: list[nn.Module] = []
+        if ni_in != nf_out:
+            id_layers.append(ConvBnAct(ni_in, nf_out, ks=1, act=False))
+        if stride != 1:
+            id_layers.insert(0, nn.AvgPool1d(2, ceil_mode=True))
+        self.idpath = nn.Sequential(*id_layers)
+        self.act = nn.ReLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.convs(x) + self.idpath(x))
+
+
+@register_model(
+    name="xresnet1d-benchmark",
+    description="XResNet1d from PTB-XL benchmark (fixed feature dim, concat pooling)",
+)
+class XResNet1dBenchmark(nn.Module):
+    """
+    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)
+
+    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
+    """
+
+    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,
+    ):
+        super().__init__()
+
+        if layers is None:
+            layers = [2, 2, 2, 2]
+
+        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)
+
+        block_szs = [64 // expansion, *([base_channels] * len(layers))]
+
+        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)
+            ]
+        )
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    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)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        return self.head(x)
+
+
+def xresnet1d18_benchmark(**kwargs) -> XResNet1dBenchmark:
+    return XResNet1dBenchmark(expansion=1, layers=[2, 2, 2, 2], **kwargs)
+
+
+def xresnet1d50_benchmark(**kwargs) -> XResNet1dBenchmark:
+    return XResNet1dBenchmark(expansion=4, layers=[3, 4, 6, 3], **kwargs)
+
+
+def xresnet1d101_benchmark(**kwargs) -> XResNet1dBenchmark:
+    return XResNet1dBenchmark(expansion=4, layers=[3, 4, 23, 3], **kwargs)