deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/models/resnet1d_wang.py

@@ -0,0 +1,182 @@
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class BasicBlock1DWang(nn.Module):
+    """
+    Basic residual block for the Wang ResNet variant.
+
+    Uses two convolutional layers with batch normalization and a skip connection.
+    The second convolution uses a smaller kernel (kernel_size // 2 + 1).
+
+    Args:
+        in_channels: Number of input channels
+        out_channels: Number of output channels
+        kernel_size: Primary kernel size (default: 5)
+        stride: Stride for the first convolution (default: 1)
+        downsample: Optional downsample module for the skip connection
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 5,
+        stride: int = 1,
+        downsample: nn.Module | None = None,
+    ):
+        super().__init__()
+
+        ks2 = kernel_size // 2 + 1
+
+        self.conv1 = nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2,
+            bias=False,
+        )
+        self.bn1 = nn.BatchNorm1d(out_channels)
+        self.conv2 = nn.Conv1d(
+            out_channels,
+            out_channels,
+            kernel_size=ks2,
+            stride=1,
+            padding=(ks2 - 1) // 2,
+            bias=False,
+        )
+        self.bn2 = nn.BatchNorm1d(out_channels)
+        self.downsample = downsample
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = func.relu(self.bn1(self.conv1(x)), inplace=True)
+        out = self.bn2(self.conv2(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        return func.relu(out, inplace=True)
+
+
+@register_model(
+    name="resnet-wang",
+    description="ResNet (Wang et al.) for time series classification",
+)
+class ResNetWang(nn.Module):
+    """
+    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)
+
+    @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.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.global_pool(x)
+        return torch.flatten(x, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        return self.classifier(x)

deepecgkit/models/se_resnet1d.py

@@ -0,0 +1,232 @@
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class SqueezeExcitation1D(nn.Module):
+    """
+    Squeeze-and-Excitation block for 1D signals.
+
+    Learns per-channel attention weights via global pooling followed by a
+    bottleneck MLP, allowing the network to emphasize informative channels
+    (e.g., specific ECG leads) and suppress less useful ones.
+
+    Args:
+        channels: Number of input/output channels
+        reduction: Channel reduction ratio for the bottleneck (default: 16)
+    """
+
+    def __init__(self, channels: int, reduction: int = 16):
+        super().__init__()
+        mid = max(channels // reduction, 1)
+        self.fc1 = nn.Linear(channels, mid, bias=False)
+        self.fc2 = nn.Linear(mid, channels, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, _ = x.shape
+        scale = x.mean(dim=2)
+        scale = func.relu(self.fc1(scale))
+        scale = torch.sigmoid(self.fc2(scale))
+        return x * scale.view(b, c, 1)
+
+
+class SEResidualBlock1D(nn.Module):
+    """
+    Residual block with Squeeze-and-Excitation attention.
+
+    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
+        se_reduction: SE reduction ratio (default: 16)
+        dropout_rate: Dropout probability (default: 0.0)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 7,
+        stride: int = 1,
+        downsample: nn.Module | None = None,
+        se_reduction: int = 16,
+        dropout_rate: float = 0.0,
+    ):
+        super().__init__()
+
+        self.conv1 = nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            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.se = SqueezeExcitation1D(out_channels, reduction=se_reduction)
+        self.downsample = downsample
+        self.dropout = nn.Dropout(dropout_rate) if dropout_rate > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+
+        out = func.relu(self.bn1(self.conv1(x)))
+        out = self.dropout(out)
+        out = self.bn2(self.conv2(out))
+        out = self.se(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = func.relu(out)
+        return out
+
+
+@register_model(
+    name="se-resnet",
+    description="ResNet with Squeeze-and-Excitation channel attention for ECG",
+)
+class SEResNet1D(nn.Module):
+    """
+    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
+    """
+
+    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):
+        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)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    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
+
+    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/simple_cnn.py

@@ -0,0 +1,82 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+@register_model(
+    name="simple-cnn",
+    description="Lightweight CNN for fast inference",
+)
+class SimpleCNN(nn.Module):
+    """
+    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.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(128, output_size),
+        )
+
+    @property
+    def feature_dim(self) -> int:
+        return 256
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.features(x)
+        return torch.flatten(x, 1)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.classifier(x)
+        return x

deepecgkit/models/tcn.py

@@ -0,0 +1,178 @@
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class CausalConv1d(nn.Module):
+    """
+    Causal 1D convolution that only looks at past and present timesteps.
+
+    Uses left-side padding to ensure no information leakage from the future,
+    making it suitable for real-time or streaming ECG inference.
+
+    Args:
+        in_channels: Number of input channels
+        out_channels: Number of output channels
+        kernel_size: Convolution kernel size
+        dilation: Dilation factor for dilated convolutions
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        self.padding = (kernel_size - 1) * dilation
+        self.conv = nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            padding=self.padding,
+            bias=False,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.conv(x)
+        if self.padding > 0:
+            out = out[:, :, : -self.padding]
+        return out
+
+
+class TemporalBlock(nn.Module):
+    """
+    Single temporal block with dilated causal convolutions and residual connection.
+
+    Each block applies two dilated causal convolutions with weight normalization,
+    ReLU activations, and dropout, plus a 1x1 shortcut if channel sizes differ.
+
+    Args:
+        in_channels: Number of input channels
+        out_channels: Number of output channels
+        kernel_size: Convolution kernel size (default: 7)
+        dilation: Dilation factor (default: 1)
+        dropout_rate: Dropout probability (default: 0.2)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 7,
+        dilation: int = 1,
+        dropout_rate: float = 0.2,
+    ):
+        super().__init__()
+
+        self.causal_conv1 = CausalConv1d(in_channels, out_channels, kernel_size, dilation)
+        nn.utils.parametrizations.weight_norm(self.causal_conv1.conv)
+        self.bn1 = nn.BatchNorm1d(out_channels)
+        self.dropout1 = nn.Dropout(dropout_rate)
+
+        self.causal_conv2 = CausalConv1d(out_channels, out_channels, kernel_size, dilation)
+        nn.utils.parametrizations.weight_norm(self.causal_conv2.conv)
+        self.bn2 = nn.BatchNorm1d(out_channels)
+        self.dropout2 = nn.Dropout(dropout_rate)
+
+        self.shortcut = (
+            nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.causal_conv1(x)
+        out = func.relu(self.bn1(out))
+        out = self.dropout1(out)
+
+        out = self.causal_conv2(out)
+        out = func.relu(self.bn2(out))
+        out = self.dropout2(out)
+
+        return func.relu(out + self.shortcut(x))
+
+
+@register_model(
+    name="tcn",
+    description="Temporal Convolutional Network with dilated causal convolutions",
+)
+class TCN(nn.Module):
+    """
+    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)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    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
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        x = self.classifier(x)
+        return x