wavedl 1.5.7__py3-none-any.whl → 1.6.1__py3-none-any.whl

wavedl/models/swin.py

@@ -1,443 +1,443 @@
-"""
-Swin Transformer: Hierarchical Vision Transformer with Shifted Windows
-=======================================================================
-
-State-of-the-art vision transformer that computes self-attention within
-local windows while enabling cross-window connections via shifting.
-Achieves excellent accuracy with linear computational complexity.
-
-**Key Innovations**:
-    - Hierarchical feature maps (like CNNs) for multi-scale processing
-    - Shifted window attention: O(n) complexity vs O(n²) for vanilla ViT
-    - Local attention with global receptive field through layer stacking
-    - Strong inductive bias for structured data
-
-**Variants**:
-    - swin_t: Tiny (28M params) - Efficient default
-    - swin_s: Small (50M params) - Better accuracy
-    - swin_b: Base (88M params) - High accuracy
-
-**Why Swin over ViT?**:
-    - Better for smaller datasets (stronger inductive bias)
-    - Handles higher resolution inputs efficiently
-    - Produces hierarchical features (useful for multi-scale patterns)
-    - More efficient memory usage
-
-**Note**: Swin Transformer is 2D-only. For 1D data, use TCN. For 3D data, use ResNet3D.
-
-References:
-    Liu, Z., et al. (2021). Swin Transformer: Hierarchical Vision Transformer
-    using Shifted Windows. ICCV 2021 (Best Paper). https://arxiv.org/abs/2103.14030
-
-Author: Ductho Le (ductho.le@outlook.com)
-"""
-
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-
-try:
-    from torchvision.models import (
-        Swin_B_Weights,
-        Swin_S_Weights,
-        Swin_T_Weights,
-        swin_b,
-        swin_s,
-        swin_t,
-    )
-
-    SWIN_AVAILABLE = True
-except ImportError:
-    SWIN_AVAILABLE = False
-
-from wavedl.models.base import BaseModel
-from wavedl.models.registry import register_model
-
-
-class SwinTransformerBase(BaseModel):
-    """
-    Base Swin Transformer class for regression tasks.
-
-    Wraps torchvision Swin Transformer with:
-    - Optional pretrained weights (ImageNet-1K or ImageNet-22K)
-    - Automatic input channel adaptation (grayscale → 3ch)
-    - Custom regression head with layer normalization
-
-    Swin Transformer excels at:
-    - Multi-scale feature extraction (dispersion curves, spectrograms)
-    - High-resolution inputs (efficient O(n) attention)
-    - Tasks requiring both local and global context
-    - Transfer learning from pretrained weights
-
-    Note: This is 2D-only. Input shape must be (H, W).
-    """
-
-    def __init__(
-        self,
-        in_shape: tuple[int, int],
-        out_size: int,
-        model_fn,
-        weights_class,
-        pretrained: bool = True,
-        dropout_rate: float = 0.3,
-        freeze_backbone: bool = False,
-        regression_hidden: int = 512,
-        **kwargs,
-    ):
-        """
-        Initialize Swin Transformer for regression.
-
-        Args:
-            in_shape: (H, W) input image dimensions
-            out_size: Number of regression output targets
-            model_fn: torchvision model constructor
-            weights_class: Pretrained weights enum class
-            pretrained: Use ImageNet pretrained weights (default: True)
-            dropout_rate: Dropout rate in regression head (default: 0.3)
-            freeze_backbone: Freeze backbone for fine-tuning (default: False)
-            regression_hidden: Hidden units in regression head (default: 512)
-        """
-        super().__init__(in_shape, out_size)
-
-        if not SWIN_AVAILABLE:
-            raise ImportError(
-                "torchvision >= 0.12 is required for Swin Transformer. "
-                "Install with: pip install torchvision>=0.12"
-            )
-
-        if len(in_shape) != 2:
-            raise ValueError(
-                f"Swin Transformer requires 2D input (H, W), got {len(in_shape)}D. "
-                "For 1D data, use TCN. For 3D data, use ResNet3D."
-            )
-
-        self.pretrained = pretrained
-        self.dropout_rate = dropout_rate
-        self.freeze_backbone = freeze_backbone
-        self.regression_hidden = regression_hidden
-
-        # Load pretrained backbone
-        weights = weights_class.IMAGENET1K_V1 if pretrained else None
-        self.backbone = model_fn(weights=weights)
-
-        # Swin Transformer head structure:
-        # head: Linear (embed_dim → num_classes)
-        # We need to get the embedding dimension from the head
-
-        in_features = self.backbone.head.in_features
-
-        # Replace head with regression head
-        # Use LayerNorm for stability (matches Transformer architecture)
-        self.backbone.head = nn.Sequential(
-            nn.LayerNorm(in_features),
-            nn.Dropout(dropout_rate),
-            nn.Linear(in_features, regression_hidden),
-            nn.GELU(),  # GELU matches Transformer's activation
-            nn.Dropout(dropout_rate * 0.5),
-            nn.Linear(regression_hidden, regression_hidden // 2),
-            nn.GELU(),
-            nn.Linear(regression_hidden // 2, out_size),
-        )
-
-        # Adapt patch embedding conv for single-channel input (3× memory savings vs expand)
-        self._adapt_input_channels()
-
-        # Optionally freeze backbone for fine-tuning (after adaptation so new conv is frozen too)
-        if freeze_backbone:
-            self._freeze_backbone()
-
-    def _adapt_input_channels(self):
-        """Modify patch embedding conv to accept single-channel input.
-
-        Instead of expanding 1→3 channels in forward (which triples memory),
-        we replace the patch embedding conv with a 1-channel version and
-        initialize weights as the mean of the pretrained RGB filters.
-        """
-        # Swin's patch embedding is at features[0][0]
-        try:
-            old_conv = self.backbone.features[0][0]
-        except (IndexError, AttributeError, TypeError) as e:
-            raise RuntimeError(
-                f"Swin patch embed structure changed in this torchvision version. "
-                f"Cannot adapt input channels. Error: {e}"
-            ) from e
-        new_conv = nn.Conv2d(
-            1,  # Single channel input
-            old_conv.out_channels,
-            kernel_size=old_conv.kernel_size,
-            stride=old_conv.stride,
-            padding=old_conv.padding,
-            dilation=old_conv.dilation,
-            groups=old_conv.groups,
-            padding_mode=old_conv.padding_mode,
-            bias=old_conv.bias is not None,
-        )
-        if self.pretrained:
-            with torch.no_grad():
-                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
-                if old_conv.bias is not None:
-                    new_conv.bias.copy_(old_conv.bias)
-        self.backbone.features[0][0] = new_conv
-
-    def _freeze_backbone(self):
-        """Freeze all backbone parameters except the head."""
-        for name, param in self.backbone.named_parameters():
-            if "head" not in name:
-                param.requires_grad = False
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass.
-
-        Args:
-            x: Input tensor of shape (B, 1, H, W)
-
-        Returns:
-            Output tensor of shape (B, out_size)
-        """
-        return self.backbone(x)
-
-    @classmethod
-    def get_default_config(cls) -> dict[str, Any]:
-        """Return default configuration for Swin Transformer."""
-        return {
-            "pretrained": True,
-            "dropout_rate": 0.3,
-            "freeze_backbone": False,
-            "regression_hidden": 512,
-        }
-
-    def get_optimizer_groups(self, base_lr: float, weight_decay: float = 0.05) -> list:
-        """
-        Get parameter groups with layer-wise learning rate decay.
-
-        Swin Transformer benefits from decaying learning rate for earlier layers.
-        This is a common practice for fine-tuning vision transformers.
-
-        Args:
-            base_lr: Base learning rate (applied to head)
-            weight_decay: Weight decay coefficient
-
-        Returns:
-            List of parameter group dictionaries
-        """
-        # Separate parameters into 4 groups for proper LR decay:
-        # 1. Head params with decay (full LR)
-        # 2. Backbone params with decay (0.1× LR)
-        # 3. Head bias/norm without decay (full LR)
-        # 4. Backbone bias/norm without decay (0.1× LR)
-        head_params = []
-        backbone_params = []
-        head_no_decay = []
-        backbone_no_decay = []
-
-        for name, param in self.backbone.named_parameters():
-            if not param.requires_grad:
-                continue
-
-            is_head = "head" in name
-            is_no_decay = "bias" in name or "norm" in name
-
-            if is_head:
-                if is_no_decay:
-                    head_no_decay.append(param)
-                else:
-                    head_params.append(param)
-            else:
-                if is_no_decay:
-                    backbone_no_decay.append(param)
-                else:
-                    backbone_params.append(param)
-
-        groups = []
-
-        if head_params:
-            groups.append(
-                {
-                    "params": head_params,
-                    "lr": base_lr,
-                    "weight_decay": weight_decay,
-                }
-            )
-
-        if backbone_params:
-            # Apply 0.1x learning rate to backbone (common for fine-tuning)
-            groups.append(
-                {
-                    "params": backbone_params,
-                    "lr": base_lr * 0.1,
-                    "weight_decay": weight_decay,
-                }
-            )
-
-        if head_no_decay:
-            groups.append(
-                {
-                    "params": head_no_decay,
-                    "lr": base_lr,
-                    "weight_decay": 0.0,
-                }
-            )
-
-        if backbone_no_decay:
-            # Backbone bias/norm also gets 0.1× LR to match intended decay
-            groups.append(
-                {
-                    "params": backbone_no_decay,
-                    "lr": base_lr * 0.1,
-                    "weight_decay": 0.0,
-                }
-            )
-
-        return groups if groups else [{"params": self.parameters(), "lr": base_lr}]
-
-
-# =============================================================================
-# REGISTERED MODEL VARIANTS
-# =============================================================================
-
-
-@register_model("swin_t")
-class SwinTiny(SwinTransformerBase):
-    """
-    Swin-T (Tiny): Efficient default for most wave-based tasks.
-
-    ~28M parameters. Good balance of accuracy and computational cost.
-    Outperforms ResNet50 while being more efficient.
-
-    Recommended for:
-        - Default choice for 2D wave data
-        - Dispersion curves, spectrograms, B-scans
-        - When hierarchical features matter
-        - Transfer learning with limited data
-
-    Architecture:
-        - Patch size: 4×4
-        - Window size: 7×7
-        - Embed dim: 96
-        - Depths: [2, 2, 6, 2]
-        - Heads: [3, 6, 12, 24]
-
-    Args:
-        in_shape: (H, W) image dimensions
-        out_size: Number of regression targets
-        pretrained: Use ImageNet pretrained weights (default: True)
-        dropout_rate: Dropout rate in head (default: 0.3)
-        freeze_backbone: Freeze backbone for fine-tuning (default: False)
-        regression_hidden: Hidden units in regression head (default: 512)
-
-    Example:
-        >>> model = SwinTiny(in_shape=(224, 224), out_size=3)
-        >>> x = torch.randn(4, 1, 224, 224)
-        >>> out = model(x)  # (4, 3)
-    """
-
-    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
-        super().__init__(
-            in_shape=in_shape,
-            out_size=out_size,
-            model_fn=swin_t,
-            weights_class=Swin_T_Weights,
-            **kwargs,
-        )
-
-    def __repr__(self) -> str:
-        pt = "pretrained" if self.pretrained else "scratch"
-        return f"Swin_Tiny({pt}, in={self.in_shape}, out={self.out_size})"
-
-
-@register_model("swin_s")
-class SwinSmall(SwinTransformerBase):
-    """
-    Swin-S (Small): Higher accuracy with moderate compute.
-
-    ~50M parameters. Better accuracy than Swin-T for larger datasets.
-
-    Recommended for:
-        - Larger datasets (>20k samples)
-        - When Swin-T doesn't provide enough capacity
-        - Complex multi-scale patterns
-
-    Architecture:
-        - Patch size: 4×4
-        - Window size: 7×7
-        - Embed dim: 96
-        - Depths: [2, 2, 18, 2]  (deeper stage 3)
-        - Heads: [3, 6, 12, 24]
-
-    Args:
-        in_shape: (H, W) image dimensions
-        out_size: Number of regression targets
-        pretrained: Use ImageNet pretrained weights (default: True)
-        dropout_rate: Dropout rate in head (default: 0.3)
-        freeze_backbone: Freeze backbone for fine-tuning (default: False)
-        regression_hidden: Hidden units in regression head (default: 512)
-
-    Example:
-        >>> model = SwinSmall(in_shape=(224, 224), out_size=3)
-        >>> x = torch.randn(4, 1, 224, 224)
-        >>> out = model(x)  # (4, 3)
-    """
-
-    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
-        super().__init__(
-            in_shape=in_shape,
-            out_size=out_size,
-            model_fn=swin_s,
-            weights_class=Swin_S_Weights,
-            **kwargs,
-        )
-
-    def __repr__(self) -> str:
-        pt = "pretrained" if self.pretrained else "scratch"
-        return f"Swin_Small({pt}, in={self.in_shape}, out={self.out_size})"
-
-
-@register_model("swin_b")
-class SwinBase(SwinTransformerBase):
-    """
-    Swin-B (Base): Maximum accuracy for large-scale tasks.
-
-    ~88M parameters. Best accuracy but requires more compute and data.
-
-    Recommended for:
-        - Very large datasets (>50k samples)
-        - When accuracy is more important than efficiency
-        - HPC environments with ample GPU memory
-        - Research experiments
-
-    Architecture:
-        - Patch size: 4×4
-        - Window size: 7×7
-        - Embed dim: 128
-        - Depths: [2, 2, 18, 2]
-        - Heads: [4, 8, 16, 32]
-
-    Args:
-        in_shape: (H, W) image dimensions
-        out_size: Number of regression targets
-        pretrained: Use ImageNet pretrained weights (default: True)
-        dropout_rate: Dropout rate in head (default: 0.3)
-        freeze_backbone: Freeze backbone for fine-tuning (default: False)
-        regression_hidden: Hidden units in regression head (default: 512)
-
-    Example:
-        >>> model = SwinBase(in_shape=(224, 224), out_size=3)
-        >>> x = torch.randn(4, 1, 224, 224)
-        >>> out = model(x)  # (4, 3)
-    """
-
-    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
-        super().__init__(
-            in_shape=in_shape,
-            out_size=out_size,
-            model_fn=swin_b,
-            weights_class=Swin_B_Weights,
-            **kwargs,
-        )
-
-    def __repr__(self) -> str:
-        pt = "pretrained" if self.pretrained else "scratch"
-        return f"Swin_Base({pt}, in={self.in_shape}, out={self.out_size})"
+"""
+Swin Transformer: Hierarchical Vision Transformer with Shifted Windows
+=======================================================================
+
+State-of-the-art vision transformer that computes self-attention within
+local windows while enabling cross-window connections via shifting.
+Achieves excellent accuracy with linear computational complexity.
+
+**Key Innovations**:
+    - Hierarchical feature maps (like CNNs) for multi-scale processing
+    - Shifted window attention: O(n) complexity vs O(n²) for vanilla ViT
+    - Local attention with global receptive field through layer stacking
+    - Strong inductive bias for structured data
+
+**Variants**:
+    - swin_t: Tiny (28M params) - Efficient default
+    - swin_s: Small (50M params) - Better accuracy
+    - swin_b: Base (88M params) - High accuracy
+
+**Why Swin over ViT?**:
+    - Better for smaller datasets (stronger inductive bias)
+    - Handles higher resolution inputs efficiently
+    - Produces hierarchical features (useful for multi-scale patterns)
+    - More efficient memory usage
+
+**Note**: Swin Transformer is 2D-only. For 1D data, use TCN. For 3D data, use ResNet3D.
+
+References:
+    Liu, Z., et al. (2021). Swin Transformer: Hierarchical Vision Transformer
+    using Shifted Windows. ICCV 2021 (Best Paper). https://arxiv.org/abs/2103.14030
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+
+try:
+    from torchvision.models import (
+        Swin_B_Weights,
+        Swin_S_Weights,
+        Swin_T_Weights,
+        swin_b,
+        swin_s,
+        swin_t,
+    )
+
+    SWIN_AVAILABLE = True
+except ImportError:
+    SWIN_AVAILABLE = False
+
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+class SwinTransformerBase(BaseModel):
+    """
+    Base Swin Transformer class for regression tasks.
+
+    Wraps torchvision Swin Transformer with:
+    - Optional pretrained weights (ImageNet-1K or ImageNet-22K)
+    - Automatic input channel adaptation (grayscale → 3ch)
+    - Custom regression head with layer normalization
+
+    Swin Transformer excels at:
+    - Multi-scale feature extraction (dispersion curves, spectrograms)
+    - High-resolution inputs (efficient O(n) attention)
+    - Tasks requiring both local and global context
+    - Transfer learning from pretrained weights
+
+    Note: This is 2D-only. Input shape must be (H, W).
+    """
+
+    def __init__(
+        self,
+        in_shape: tuple[int, int],
+        out_size: int,
+        model_fn,
+        weights_class,
+        pretrained: bool = True,
+        dropout_rate: float = 0.3,
+        freeze_backbone: bool = False,
+        regression_hidden: int = 512,
+        **kwargs,
+    ):
+        """
+        Initialize Swin Transformer for regression.
+
+        Args:
+            in_shape: (H, W) input image dimensions
+            out_size: Number of regression output targets
+            model_fn: torchvision model constructor
+            weights_class: Pretrained weights enum class
+            pretrained: Use ImageNet pretrained weights (default: True)
+            dropout_rate: Dropout rate in regression head (default: 0.3)
+            freeze_backbone: Freeze backbone for fine-tuning (default: False)
+            regression_hidden: Hidden units in regression head (default: 512)
+        """
+        super().__init__(in_shape, out_size)
+
+        if not SWIN_AVAILABLE:
+            raise ImportError(
+                "torchvision >= 0.12 is required for Swin Transformer. "
+                "Install with: pip install torchvision>=0.12"
+            )
+
+        if len(in_shape) != 2:
+            raise ValueError(
+                f"Swin Transformer requires 2D input (H, W), got {len(in_shape)}D. "
+                "For 1D data, use TCN. For 3D data, use ResNet3D."
+            )
+
+        self.pretrained = pretrained
+        self.dropout_rate = dropout_rate
+        self.freeze_backbone = freeze_backbone
+        self.regression_hidden = regression_hidden
+
+        # Load pretrained backbone
+        weights = weights_class.IMAGENET1K_V1 if pretrained else None
+        self.backbone = model_fn(weights=weights)
+
+        # Swin Transformer head structure:
+        # head: Linear (embed_dim → num_classes)
+        # We need to get the embedding dimension from the head
+
+        in_features = self.backbone.head.in_features
+
+        # Replace head with regression head
+        # Use LayerNorm for stability (matches Transformer architecture)
+        self.backbone.head = nn.Sequential(
+            nn.LayerNorm(in_features),
+            nn.Dropout(dropout_rate),
+            nn.Linear(in_features, regression_hidden),
+            nn.GELU(),  # GELU matches Transformer's activation
+            nn.Dropout(dropout_rate * 0.5),
+            nn.Linear(regression_hidden, regression_hidden // 2),
+            nn.GELU(),
+            nn.Linear(regression_hidden // 2, out_size),
+        )
+
+        # Adapt patch embedding conv for single-channel input (3× memory savings vs expand)
+        self._adapt_input_channels()
+
+        # Optionally freeze backbone for fine-tuning (after adaptation so new conv is frozen too)
+        if freeze_backbone:
+            self._freeze_backbone()
+
+    def _adapt_input_channels(self):
+        """Modify patch embedding conv to accept single-channel input.
+
+        Instead of expanding 1→3 channels in forward (which triples memory),
+        we replace the patch embedding conv with a 1-channel version and
+        initialize weights as the mean of the pretrained RGB filters.
+        """
+        # Swin's patch embedding is at features[0][0]
+        try:
+            old_conv = self.backbone.features[0][0]
+        except (IndexError, AttributeError, TypeError) as e:
+            raise RuntimeError(
+                f"Swin patch embed structure changed in this torchvision version. "
+                f"Cannot adapt input channels. Error: {e}"
+            ) from e
+        new_conv = nn.Conv2d(
+            1,  # Single channel input
+            old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+            dilation=old_conv.dilation,
+            groups=old_conv.groups,
+            padding_mode=old_conv.padding_mode,
+            bias=old_conv.bias is not None,
+        )
+        if self.pretrained:
+            with torch.no_grad():
+                new_conv.weight.copy_(old_conv.weight.mean(dim=1, keepdim=True))
+                if old_conv.bias is not None:
+                    new_conv.bias.copy_(old_conv.bias)
+        self.backbone.features[0][0] = new_conv
+
+    def _freeze_backbone(self):
+        """Freeze all backbone parameters except the head."""
+        for name, param in self.backbone.named_parameters():
+            if "head" not in name:
+                param.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Args:
+            x: Input tensor of shape (B, 1, H, W)
+
+        Returns:
+            Output tensor of shape (B, out_size)
+        """
+        return self.backbone(x)
+
+    @classmethod
+    def get_default_config(cls) -> dict[str, Any]:
+        """Return default configuration for Swin Transformer."""
+        return {
+            "pretrained": True,
+            "dropout_rate": 0.3,
+            "freeze_backbone": False,
+            "regression_hidden": 512,
+        }
+
+    def get_optimizer_groups(self, base_lr: float, weight_decay: float = 0.05) -> list:
+        """
+        Get parameter groups with layer-wise learning rate decay.
+
+        Swin Transformer benefits from decaying learning rate for earlier layers.
+        This is a common practice for fine-tuning vision transformers.
+
+        Args:
+            base_lr: Base learning rate (applied to head)
+            weight_decay: Weight decay coefficient
+
+        Returns:
+            List of parameter group dictionaries
+        """
+        # Separate parameters into 4 groups for proper LR decay:
+        # 1. Head params with decay (full LR)
+        # 2. Backbone params with decay (0.1× LR)
+        # 3. Head bias/norm without decay (full LR)
+        # 4. Backbone bias/norm without decay (0.1× LR)
+        head_params = []
+        backbone_params = []
+        head_no_decay = []
+        backbone_no_decay = []
+
+        for name, param in self.backbone.named_parameters():
+            if not param.requires_grad:
+                continue
+
+            is_head = "head" in name
+            is_no_decay = "bias" in name or "norm" in name
+
+            if is_head:
+                if is_no_decay:
+                    head_no_decay.append(param)
+                else:
+                    head_params.append(param)
+            else:
+                if is_no_decay:
+                    backbone_no_decay.append(param)
+                else:
+                    backbone_params.append(param)
+
+        groups = []
+
+        if head_params:
+            groups.append(
+                {
+                    "params": head_params,
+                    "lr": base_lr,
+                    "weight_decay": weight_decay,
+                }
+            )
+
+        if backbone_params:
+            # Apply 0.1x learning rate to backbone (common for fine-tuning)
+            groups.append(
+                {
+                    "params": backbone_params,
+                    "lr": base_lr * 0.1,
+                    "weight_decay": weight_decay,
+                }
+            )
+
+        if head_no_decay:
+            groups.append(
+                {
+                    "params": head_no_decay,
+                    "lr": base_lr,
+                    "weight_decay": 0.0,
+                }
+            )
+
+        if backbone_no_decay:
+            # Backbone bias/norm also gets 0.1× LR to match intended decay
+            groups.append(
+                {
+                    "params": backbone_no_decay,
+                    "lr": base_lr * 0.1,
+                    "weight_decay": 0.0,
+                }
+            )
+
+        return groups if groups else [{"params": self.parameters(), "lr": base_lr}]
+
+
+# =============================================================================
+# REGISTERED MODEL VARIANTS
+# =============================================================================
+
+
+@register_model("swin_t")
+class SwinTiny(SwinTransformerBase):
+    """
+    Swin-T (Tiny): Efficient default for most wave-based tasks.
+
+    ~27.5M backbone parameters. Good balance of accuracy and computational cost.
+    Outperforms ResNet50 while being more efficient.
+
+    Recommended for:
+        - Default choice for 2D wave data
+        - Dispersion curves, spectrograms, B-scans
+        - When hierarchical features matter
+        - Transfer learning with limited data
+
+    Architecture:
+        - Patch size: 4×4
+        - Window size: 7×7
+        - Embed dim: 96
+        - Depths: [2, 2, 6, 2]
+        - Heads: [3, 6, 12, 24]
+
+    Args:
+        in_shape: (H, W) image dimensions
+        out_size: Number of regression targets
+        pretrained: Use ImageNet pretrained weights (default: True)
+        dropout_rate: Dropout rate in head (default: 0.3)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 512)
+
+    Example:
+        >>> model = SwinTiny(in_shape=(224, 224), out_size=3)
+        >>> x = torch.randn(4, 1, 224, 224)
+        >>> out = model(x)  # (4, 3)
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_fn=swin_t,
+            weights_class=Swin_T_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"Swin_Tiny({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("swin_s")
+class SwinSmall(SwinTransformerBase):
+    """
+    Swin-S (Small): Higher accuracy with moderate compute.
+
+    ~48.8M backbone parameters. Better accuracy than Swin-T for larger datasets.
+
+    Recommended for:
+        - Larger datasets (>20k samples)
+        - When Swin-T doesn't provide enough capacity
+        - Complex multi-scale patterns
+
+    Architecture:
+        - Patch size: 4×4
+        - Window size: 7×7
+        - Embed dim: 96
+        - Depths: [2, 2, 18, 2]  (deeper stage 3)
+        - Heads: [3, 6, 12, 24]
+
+    Args:
+        in_shape: (H, W) image dimensions
+        out_size: Number of regression targets
+        pretrained: Use ImageNet pretrained weights (default: True)
+        dropout_rate: Dropout rate in head (default: 0.3)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 512)
+
+    Example:
+        >>> model = SwinSmall(in_shape=(224, 224), out_size=3)
+        >>> x = torch.randn(4, 1, 224, 224)
+        >>> out = model(x)  # (4, 3)
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_fn=swin_s,
+            weights_class=Swin_S_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"Swin_Small({pt}, in={self.in_shape}, out={self.out_size})"
+
+
+@register_model("swin_b")
+class SwinBase(SwinTransformerBase):
+    """
+    Swin-B (Base): Maximum accuracy for large-scale tasks.
+
+    ~86.7M backbone parameters. Best accuracy but requires more compute and data.
+
+    Recommended for:
+        - Very large datasets (>50k samples)
+        - When accuracy is more important than efficiency
+        - HPC environments with ample GPU memory
+        - Research experiments
+
+    Architecture:
+        - Patch size: 4×4
+        - Window size: 7×7
+        - Embed dim: 128
+        - Depths: [2, 2, 18, 2]
+        - Heads: [4, 8, 16, 32]
+
+    Args:
+        in_shape: (H, W) image dimensions
+        out_size: Number of regression targets
+        pretrained: Use ImageNet pretrained weights (default: True)
+        dropout_rate: Dropout rate in head (default: 0.3)
+        freeze_backbone: Freeze backbone for fine-tuning (default: False)
+        regression_hidden: Hidden units in regression head (default: 512)
+
+    Example:
+        >>> model = SwinBase(in_shape=(224, 224), out_size=3)
+        >>> x = torch.randn(4, 1, 224, 224)
+        >>> out = model(x)  # (4, 3)
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_fn=swin_b,
+            weights_class=Swin_B_Weights,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        pt = "pretrained" if self.pretrained else "scratch"
+        return f"Swin_Base({pt}, in={self.in_shape}, out={self.out_size})"