deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/models/mamba1d.py

@@ -0,0 +1,232 @@
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class SelectiveSSM(nn.Module):
+    """
+    Selective State Space Model implemented in pure PyTorch.
+
+    Implements a discretized state space: h_t = A_bar * h_{t-1} + B_bar * x_t,
+    y_t = C * h_t, where B, C, and delta are input-dependent (selective).
+
+    Args:
+        d_inner: Inner model dimension
+        d_state: State space dimension (default: 16)
+        d_conv: Local convolution width (default: 4)
+    """
+
+    def __init__(self, d_inner: int, d_state: int = 16, d_conv: int = 4):
+        super().__init__()
+        self.d_inner = d_inner
+        self.d_state = d_state
+
+        self.conv1d = nn.Conv1d(
+            d_inner,
+            d_inner,
+            kernel_size=d_conv,
+            groups=d_inner,
+            padding=d_conv - 1,
+        )
+
+        self.x_proj = nn.Linear(d_inner, d_state * 2 + 1, bias=False)
+        self.dt_proj = nn.Linear(1, d_inner, bias=True)
+
+        a_log = torch.log(torch.arange(1, d_state + 1, dtype=torch.float32))
+        a_log = a_log.unsqueeze(0).expand(d_inner, -1)
+        self.a_log = nn.Parameter(a_log)
+
+        self.D = nn.Parameter(torch.ones(d_inner))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, seq_len, _ = x.shape
+
+        x_conv = x.transpose(1, 2)
+        x_conv = self.conv1d(x_conv)[:, :, :seq_len]
+        x_conv = func.silu(x_conv).transpose(1, 2)
+
+        proj = self.x_proj(x_conv)
+        b_sel = proj[:, :, : self.d_state]
+        c_sel = proj[:, :, self.d_state : self.d_state * 2]
+        delta = proj[:, :, self.d_state * 2 :]
+
+        delta = func.softplus(self.dt_proj(delta))
+
+        a_coeff = -torch.exp(self.a_log)
+
+        outputs = []
+        h = torch.zeros(batch, self.d_inner, self.d_state, device=x.device, dtype=x.dtype)
+
+        for t in range(seq_len):
+            dt = delta[:, t].unsqueeze(-1)
+            a_bar = torch.exp(dt * a_coeff)
+            b_bar = dt * b_sel[:, t].unsqueeze(1)
+            x_t = x_conv[:, t].unsqueeze(-1)
+
+            h = a_bar * h + b_bar * x_t
+            y_t = (h * c_sel[:, t].unsqueeze(1)).sum(-1)
+            outputs.append(y_t)
+
+        y = torch.stack(outputs, dim=1)
+        return y + x * self.D
+
+
+class MambaBlock(nn.Module):
+    """
+    Single Mamba block: norm -> linear expand -> SSM -> gated output -> project + residual.
+
+    Args:
+        d_model: Model dimension
+        d_state: State space dimension (default: 16)
+        d_conv: Local convolution width (default: 4)
+        expansion_factor: Inner dimension expansion (default: 2)
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_state: int = 16,
+        d_conv: int = 4,
+        expansion_factor: int = 2,
+    ):
+        super().__init__()
+        d_inner = d_model * expansion_factor
+
+        self.norm = nn.LayerNorm(d_model)
+        self.in_proj = nn.Linear(d_model, d_inner * 2)
+        self.ssm = SelectiveSSM(d_inner, d_state, d_conv)
+        self.out_proj = nn.Linear(d_inner, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.norm(x)
+        x, z = self.in_proj(x).chunk(2, dim=-1)
+        x = self.ssm(x)
+        x = x * func.silu(z)
+        x = self.out_proj(x)
+        return x + residual
+
+
+class BidirectionalMamba(nn.Module):
+    """
+    Bidirectional Mamba: runs forward and backward SSMs and combines outputs.
+
+    Args:
+        d_model: Model dimension
+        d_state: State space dimension (default: 16)
+        d_conv: Local convolution width (default: 4)
+        expansion_factor: Inner dimension expansion (default: 2)
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_state: int = 16,
+        d_conv: int = 4,
+        expansion_factor: int = 2,
+    ):
+        super().__init__()
+        self.forward_block = MambaBlock(d_model, d_state, d_conv, expansion_factor)
+        self.backward_block = MambaBlock(d_model, d_state, d_conv, expansion_factor)
+        self.combine = nn.Linear(d_model * 2, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        forward_out = self.forward_block(x)
+        backward_out = self.backward_block(x.flip(1)).flip(1)
+        return self.combine(torch.cat([forward_out, backward_out], dim=-1))
+
+
+@register_model(
+    name="mamba",
+    description="Bidirectional Mamba state space model with linear complexity for ECG",
+)
+class Mamba1D(nn.Module):
+    """
+    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)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    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
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        x = self.classifier(x)
+        return x

deepecgkit/models/medformer.py

@@ -0,0 +1,229 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class PatchEmbedding1D(nn.Module):
+    """
+    Converts 1D signal into patch embeddings at a given patch size.
+
+    Args:
+        input_channels: Number of input ECG channels
+        d_model: Output embedding dimension
+        patch_size: Size of each non-overlapping patch
+    """
+
+    def __init__(self, input_channels: int, d_model: int, patch_size: int):
+        super().__init__()
+        self.proj = nn.Conv1d(input_channels, d_model, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        return x.transpose(1, 2)
+
+
+class CrossChannelAttention(nn.Module):
+    """
+    Cross-channel attention for multi-lead ECG signals.
+
+    For single-lead ECG this acts as an identity-like operation.
+
+    Args:
+        d_model: Embedding dimension
+        nhead: Number of attention heads (default: 4)
+        dropout: Dropout probability (default: 0.1)
+    """
+
+    def __init__(self, d_model: int, nhead: int = 4, dropout: float = 0.1):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x, _ = self.attn(x, x, x)
+        return self.norm(x + residual)
+
+
+class IntraGranularityBlock(nn.Module):
+    """
+    Self-attention within patches of one granularity level.
+
+    Args:
+        d_model: Embedding dimension
+        nhead: Number of attention heads (default: 8)
+        dim_feedforward: Feedforward dimension (default: 256)
+        dropout: Dropout probability (default: 0.1)
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int = 8,
+        dim_feedforward: int = 256,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encoder(x)
+
+
+class InterGranularityAttention(nn.Module):
+    """
+    Cross-attention to fuse information across different granularities.
+
+    Pools each granularity to a single vector, applies self-attention
+    across granularities, and returns updated representations.
+
+    Args:
+        d_model: Embedding dimension
+        num_granularities: Number of granularity levels
+        nhead: Number of attention heads (default: 4)
+        dropout: Dropout probability (default: 0.1)
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_granularities: int,
+        nhead: int = 4,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.norm = nn.LayerNorm(d_model)
+        self.num_granularities = num_granularities
+
+    def forward(self, granularity_features: list[torch.Tensor]) -> list[torch.Tensor]:
+        pooled = [f.mean(dim=1) for f in granularity_features]
+        stacked = torch.stack(pooled, dim=1)
+        attn_out, _ = self.attn(stacked, stacked, stacked)
+        attn_out = self.norm(attn_out + stacked)
+
+        updated = []
+        for i, feat in enumerate(granularity_features):
+            scale = attn_out[:, i].unsqueeze(1)
+            updated.append(feat + scale)
+        return updated
+
+
+@register_model(
+    name="medformer",
+    description="Multi-granularity patching Transformer for medical time series (NeurIPS 2024)",
+)
+class Medformer(nn.Module):
+    """
+    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)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        granularity_features = []
+        for patch_emb, pos_emb in zip(self.patch_embeddings, self.pos_embeddings):
+            patches = patch_emb(x)
+            num_patches = patches.size(1)
+            patches = patches + pos_emb[:, :num_patches]
+            granularity_features.append(patches)
+
+        for intra_layer, inter_layer in zip(self.intra_blocks, self.inter_blocks):
+            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)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        x = self.classifier(x)
+        return x

deepecgkit/models/resnet1d.py

@@ -0,0 +1,189 @@
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class ResidualBlock1D(nn.Module):
+    """
+    1D Residual block for ECG signal processing.
+
+    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
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 7,
+        stride: int = 1,
+        downsample=None,
+    ):
+        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.downsample = downsample
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = func.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = func.relu(out)
+
+        return out
+
+
+@register_model(
+    name="resnet",
+    description="1D ResNet architecture adapted for ECG",
+)
+class ResNet1D(nn.Module):
+    """
+    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
+    """
+
+    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):
+        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)
+
+    @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):
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        x = self.fc(x)
+        return x