wavedl 1.2.0__py3-none-any.whl

wavedl/models/unet.py

@@ -0,0 +1,304 @@
+"""
+U-Net: Encoder-Decoder Architecture for Regression
+====================================================
+
+A dimension-agnostic U-Net implementation for tasks requiring either:
+- Spatial output (e.g., velocity field prediction)
+- Vector output (global pooling → regression head)
+
+**Dimensionality Support**:
+    - 1D: Waveforms, signals (N, 1, L) → Conv1d
+    - 2D: Images, spectrograms (N, 1, H, W) → Conv2d
+    - 3D: Volumetric data (N, 1, D, H, W) → Conv3d
+
+**Variants**:
+    - unet: Full encoder-decoder with spatial output capability
+    - unet_regression: U-Net with global pooling for vector regression
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+# Type alias for spatial shapes
+SpatialShape = tuple[int] | tuple[int, int] | tuple[int, int, int]
+
+
+def _get_layers(dim: int):
+    """Get dimension-appropriate layer classes."""
+    if dim == 1:
+        return nn.Conv1d, nn.ConvTranspose1d, nn.MaxPool1d, nn.AdaptiveAvgPool1d
+    elif dim == 2:
+        return nn.Conv2d, nn.ConvTranspose2d, nn.MaxPool2d, nn.AdaptiveAvgPool2d
+    elif dim == 3:
+        return nn.Conv3d, nn.ConvTranspose3d, nn.MaxPool3d, nn.AdaptiveAvgPool3d
+    else:
+        raise ValueError(f"Unsupported dimensionality: {dim}D")
+
+
+class DoubleConv(nn.Module):
+    """Double convolution block: Conv-GN-ReLU-Conv-GN-ReLU."""
+
+    def __init__(self, in_channels: int, out_channels: int, dim: int = 2):
+        super().__init__()
+        Conv = _get_layers(dim)[0]
+
+        num_groups = min(32, out_channels)
+        while out_channels % num_groups != 0 and num_groups > 1:
+            num_groups -= 1
+
+        self.double_conv = nn.Sequential(
+            Conv(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.GroupNorm(num_groups, out_channels),
+            nn.ReLU(inplace=True),
+            Conv(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.GroupNorm(num_groups, out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.double_conv(x)
+
+
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv."""
+
+    def __init__(self, in_channels: int, out_channels: int, dim: int = 2):
+        super().__init__()
+        _, _, MaxPool, _ = _get_layers(dim)
+
+        self.maxpool_conv = nn.Sequential(
+            MaxPool(2), DoubleConv(in_channels, out_channels, dim)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """Upscaling then double conv."""
+
+    def __init__(self, in_channels: int, out_channels: int, dim: int = 2):
+        super().__init__()
+        _, ConvTranspose, _, _ = _get_layers(dim)
+
+        # in_channels comes from previous layer
+        # After upconv: in_channels // 2
+        # After concat with skip (out_channels): in_channels // 2 + out_channels = in_channels
+        # Then DoubleConv: in_channels -> out_channels
+        self.up = ConvTranspose(in_channels, in_channels // 2, kernel_size=2, stride=2)
+        self.conv = DoubleConv(in_channels, out_channels, dim)
+
+    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
+        x1 = self.up(x1)
+
+        # Handle size mismatch (pad x1 to match x2)
+        if x1.shape[2:] != x2.shape[2:]:
+            diff = [x2.size(i + 2) - x1.size(i + 2) for i in range(len(x1.shape) - 2)]
+            # Pad x1 to match x2
+            pad = []
+            for d in reversed(diff):
+                pad.extend([d // 2, d - d // 2])
+            x1 = nn.functional.pad(x1, pad)
+
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class UNetBase(BaseModel):
+    """
+    Base U-Net class for regression tasks.
+
+    Standard U-Net architecture:
+    - Encoder path with downsampling
+    - Decoder path with upsampling and skip connections
+    - Optional spatial or vector output
+    """
+
+    def __init__(
+        self,
+        in_shape: SpatialShape,
+        out_size: int,
+        base_channels: int = 64,
+        depth: int = 4,
+        dropout_rate: float = 0.1,
+        spatial_output: bool = False,
+        **kwargs,
+    ):
+        super().__init__(in_shape, out_size)
+
+        self.dim = len(in_shape)
+        self.base_channels = base_channels
+        self.depth = depth
+        self.dropout_rate = dropout_rate
+        self.spatial_output = spatial_output
+
+        Conv, _, _, AdaptivePool = _get_layers(self.dim)
+
+        # Channel progression: 64 -> 128 -> 256 -> 512 (for depth=4)
+        # features[i] = base_channels * 2^i
+        features = [base_channels * (2**i) for i in range(depth + 1)]
+
+        # Initial double conv (1 -> features[0])
+        self.inc = DoubleConv(1, features[0], self.dim)
+
+        # Encoder (down path)
+        self.downs = nn.ModuleList()
+        for i in range(depth):
+            self.downs.append(Down(features[i], features[i + 1], self.dim))
+
+        # Decoder (up path)
+        self.ups = nn.ModuleList()
+        for i in range(depth):
+            # Input: features[depth - i], Skip: features[depth - 1 - i], Output: features[depth - 1 - i]
+            self.ups.append(Up(features[depth - i], features[depth - 1 - i], self.dim))
+
+        if spatial_output:
+            # Spatial output: 1x1 conv to out_size channels
+            self.outc = Conv(features[0], out_size, kernel_size=1)
+        else:
+            # Vector output: global pooling + regression head
+            self.global_pool = AdaptivePool(1)
+            self.head = nn.Sequential(
+                nn.Dropout(dropout_rate),
+                nn.Linear(features[0], 256),
+                nn.ReLU(inplace=True),
+                nn.Dropout(dropout_rate),
+                nn.Linear(256, out_size),
+            )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize weights."""
+        for m in self.modules():
+            if isinstance(
+                m,
+                (
+                    nn.Conv1d,
+                    nn.Conv2d,
+                    nn.Conv3d,
+                    nn.ConvTranspose1d,
+                    nn.ConvTranspose2d,
+                    nn.ConvTranspose3d,
+                ),
+            ):
+                nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.GroupNorm):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        # Encoder path - collect skip connections
+        x1 = self.inc(x)
+
+        skips = [x1]
+        x = x1
+        for down in self.downs:
+            x = down(x)
+            skips.append(x)
+
+        # Remove last (bottleneck output, not a skip)
+        skips = skips[:-1]
+
+        # Decoder path - use skips in reverse order
+        for up, skip in zip(self.ups, reversed(skips)):
+            x = up(x, skip)
+
+        if self.spatial_output:
+            return self.outc(x)
+        else:
+            x = self.global_pool(x)
+            x = x.flatten(1)
+            return self.head(x)
+
+    @classmethod
+    def get_default_config(cls) -> dict[str, Any]:
+        """Return default configuration."""
+        return {"base_channels": 64, "depth": 4, "dropout_rate": 0.1}
+
+
+# =============================================================================
+# REGISTERED MODEL VARIANTS
+# =============================================================================
+
+
+@register_model("unet")
+class UNet(UNetBase):
+    """
+    U-Net with spatial output capability.
+
+    Good for: Pixel/voxel-wise regression (velocity fields, spatial maps).
+
+    Note: For spatial output, out_size is the number of output channels.
+    Output shape: (B, out_size, *spatial_dims) for spatial_output=True.
+
+    Args:
+        in_shape: (L,), (H, W), or (D, H, W)
+        out_size: Number of output channels for spatial output
+        base_channels: Base channel count (default: 64)
+        depth: Number of encoder/decoder levels (default: 4)
+        spatial_output: If True, output spatial map; if False, output vector
+        dropout_rate: Dropout rate (default: 0.1)
+    """
+
+    def __init__(
+        self,
+        in_shape: SpatialShape,
+        out_size: int,
+        spatial_output: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            spatial_output=spatial_output,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        mode = "spatial" if self.spatial_output else "vector"
+        return f"UNet({self.dim}D, {mode}, in_shape={self.in_shape}, out_size={self.out_size})"
+
+
+@register_model("unet_regression")
+class UNetRegression(UNetBase):
+    """
+    U-Net for vector regression output.
+
+    Uses U-Net encoder-decoder but applies global pooling at the end
+    for standard vector regression output.
+
+    Good for: Leveraging U-Net features (multi-scale, skip connections)
+    for standard regression tasks.
+
+    Args:
+        in_shape: (L,), (H, W), or (D, H, W)
+        out_size: Number of regression targets
+        base_channels: Base channel count (default: 64)
+        depth: Number of encoder/decoder levels (default: 4)
+        dropout_rate: Dropout rate (default: 0.1)
+    """
+
+    def __init__(self, in_shape: SpatialShape, out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape, out_size=out_size, spatial_output=False, **kwargs
+        )
+
+    def __repr__(self) -> str:
+        return f"UNet_Regression({self.dim}D, in_shape={self.in_shape}, out_size={self.out_size})"

wavedl/models/vit.py

@@ -0,0 +1,372 @@
+"""
+Vision Transformer (ViT): Transformer-Based Architecture for Regression
+========================================================================
+
+A flexible Vision Transformer implementation for regression tasks.
+Supports both 1D (signals) and 2D (images) inputs via configurable patch embedding.
+
+**Dimensionality Support**:
+    - 1D: Waveforms/signals → patches are segments
+    - 2D: Images/spectrograms → patches are grid squares
+
+**Variants**:
+    - vit_tiny: Smallest (embed_dim=192, depth=12, heads=3)
+    - vit_small: Light (embed_dim=384, depth=12, heads=6)
+    - vit_base: Standard (embed_dim=768, depth=12, heads=12)
+
+Author: Ductho Le (ductho.le@outlook.com)
+Version: 1.0.0
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+# Type alias for spatial shapes
+SpatialShape = tuple[int] | tuple[int, int]
+
+
+class PatchEmbed(nn.Module):
+    """
+    Patch Embedding module that converts input into sequence of patch embeddings.
+
+    Supports 1D and 2D inputs:
+    - 1D: Input (B, 1, L) → (B, num_patches, embed_dim)
+    - 2D: Input (B, 1, H, W) → (B, num_patches, embed_dim)
+    """
+
+    def __init__(self, in_shape: SpatialShape, patch_size: int, embed_dim: int):
+        super().__init__()
+
+        self.dim = len(in_shape)
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+
+        if self.dim == 1:
+            # 1D: segment patches
+            L = in_shape[0]
+            self.num_patches = L // patch_size
+            self.proj = nn.Conv1d(
+                1, embed_dim, kernel_size=patch_size, stride=patch_size
+            )
+        elif self.dim == 2:
+            # 2D: grid patches
+            H, W = in_shape
+            self.num_patches = (H // patch_size) * (W // patch_size)
+            self.proj = nn.Conv2d(
+                1, embed_dim, kernel_size=patch_size, stride=patch_size
+            )
+        else:
+            raise ValueError(f"ViT supports 1D and 2D inputs, got {self.dim}D")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input tensor (B, 1, ..spatial..)
+
+        Returns:
+            Patch embeddings (B, num_patches, embed_dim)
+        """
+        x = self.proj(x)  # (B, embed_dim, ..reduced_spatial..)
+        x = x.flatten(2)  # (B, embed_dim, num_patches)
+        x = x.transpose(1, 2)  # (B, num_patches, embed_dim)
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    """Multi-head self-attention mechanism."""
+
+    def __init__(self, embed_dim: int, num_heads: int, dropout: float = 0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.scale = self.head_dim**-0.5
+
+        self.qkv = nn.Linear(embed_dim, embed_dim * 3, bias=True)
+        self.proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
+        qkv = qkv.permute(2, 0, 3, 1, 4)  # (3, B, heads, N, head_dim)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.dropout(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.dropout(x)
+        return x
+
+
+class MLP(nn.Module):
+    """MLP block with GELU activation."""
+
+    def __init__(self, embed_dim: int, mlp_ratio: float = 4.0, dropout: float = 0.0):
+        super().__init__()
+        hidden_dim = int(embed_dim * mlp_ratio)
+        self.fc1 = nn.Linear(embed_dim, hidden_dim)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_dim, embed_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    """Transformer encoder block with pre-norm architecture."""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.attn = MultiHeadAttention(embed_dim, num_heads, dropout)
+        self.norm2 = nn.LayerNorm(embed_dim)
+        self.mlp = MLP(embed_dim, mlp_ratio, dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class ViTBase(BaseModel):
+    """
+    Vision Transformer base class for regression.
+
+    Architecture:
+    1. Patch embedding
+    2. Add learnable position embeddings + CLS token
+    3. Transformer encoder blocks
+    4. Extract CLS token
+    5. Regression head
+    """
+
+    def __init__(
+        self,
+        in_shape: SpatialShape,
+        out_size: int,
+        patch_size: int = 16,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        dropout_rate: float = 0.1,
+        **kwargs,
+    ):
+        super().__init__(in_shape, out_size)
+
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.dropout_rate = dropout_rate
+        self.dim = len(in_shape)
+
+        # Patch embedding
+        self.patch_embed = PatchEmbed(in_shape, patch_size, embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        # Learnable CLS token and position embeddings
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.dropout = nn.Dropout(dropout_rate)
+
+        # Transformer encoder
+        self.blocks = nn.ModuleList(
+            [
+                TransformerBlock(embed_dim, num_heads, mlp_ratio, dropout_rate)
+                for _ in range(depth)
+            ]
+        )
+
+        self.norm = nn.LayerNorm(embed_dim)
+
+        # Regression head
+        self.head = nn.Sequential(
+            nn.Dropout(dropout_rate),
+            nn.Linear(embed_dim, 256),
+            nn.GELU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(256, out_size),
+        )
+
+        # Initialize weights
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize weights."""
+        # Initialize position embeddings
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+
+        # Initialize linear layers
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.LayerNorm):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x: Input tensor (B, 1, ..spatial..)
+
+        Returns:
+            Regression output (B, out_size)
+        """
+        B = x.size(0)
+
+        # Patch embedding
+        x = self.patch_embed(x)  # (B, num_patches, embed_dim)
+
+        # Prepend CLS token
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat([cls_tokens, x], dim=1)  # (B, num_patches+1, embed_dim)
+
+        # Add position embeddings
+        x = x + self.pos_embed
+        x = self.dropout(x)
+
+        # Transformer blocks
+        for block in self.blocks:
+            x = block(x)
+
+        x = self.norm(x)
+
+        # Extract CLS token for regression
+        cls_output = x[:, 0]
+
+        return self.head(cls_output)
+
+    @classmethod
+    def get_default_config(cls) -> dict[str, Any]:
+        """Return default configuration."""
+        return {"patch_size": 16, "dropout_rate": 0.1}
+
+
+# =============================================================================
+# REGISTERED MODEL VARIANTS
+# =============================================================================
+
+
+@register_model("vit_tiny")
+class ViTTiny(ViTBase):
+    """
+    ViT-Tiny: Smallest Vision Transformer variant.
+
+    ~5.7M parameters. Good for: Quick experiments, smaller datasets.
+
+    Args:
+        in_shape: (L,) for 1D or (H, W) for 2D
+        out_size: Number of regression targets
+        patch_size: Patch size (default: 16)
+        dropout_rate: Dropout rate (default: 0.1)
+    """
+
+    def __init__(
+        self, in_shape: SpatialShape, out_size: int, patch_size: int = 16, **kwargs
+    ):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            patch_size=patch_size,
+            embed_dim=192,
+            depth=12,
+            num_heads=3,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"ViT_Tiny({self.dim}D, in_shape={self.in_shape}, out_size={self.out_size})"
+        )
+
+
+@register_model("vit_small")
+class ViTSmall(ViTBase):
+    """
+    ViT-Small: Light Vision Transformer variant.
+
+    ~22M parameters. Good for: Balanced performance.
+
+    Args:
+        in_shape: (L,) for 1D or (H, W) for 2D
+        out_size: Number of regression targets
+        patch_size: Patch size (default: 16)
+        dropout_rate: Dropout rate (default: 0.1)
+    """
+
+    def __init__(
+        self, in_shape: SpatialShape, out_size: int, patch_size: int = 16, **kwargs
+    ):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            patch_size=patch_size,
+            embed_dim=384,
+            depth=12,
+            num_heads=6,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return f"ViT_Small({self.dim}D, in_shape={self.in_shape}, out_size={self.out_size})"
+
+
+@register_model("vit_base")
+class ViTBase_(ViTBase):
+    """
+    ViT-Base: Standard Vision Transformer variant.
+
+    ~86M parameters. Good for: High accuracy, larger datasets.
+
+    Args:
+        in_shape: (L,) for 1D or (H, W) for 2D
+        out_size: Number of regression targets
+        patch_size: Patch size (default: 16)
+        dropout_rate: Dropout rate (default: 0.1)
+    """
+
+    def __init__(
+        self, in_shape: SpatialShape, out_size: int, patch_size: int = 16, **kwargs
+    ):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            patch_size=patch_size,
+            embed_dim=768,
+            depth=12,
+            num_heads=12,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"ViT_Base({self.dim}D, in_shape={self.in_shape}, out_size={self.out_size})"
+        )