wavedl 1.5.6__py3-none-any.whl → 1.6.0__py3-none-any.whl

wavedl/models/fastvit.py

@@ -0,0 +1,285 @@
+"""
+FastViT: A Fast Hybrid Vision Transformer
+==========================================
+
+FastViT from Apple uses RepMixer for efficient token mixing with structural
+reparameterization - train with skip connections, deploy without.
+
+**Key Features**:
+    - RepMixer: Reparameterizable token mixing
+    - Train-time overparameterization
+    - Faster than EfficientNet/ConvNeXt on mobile
+    - CoreML compatible
+
+**Variants**:
+    - fastvit_t8: 4M params (fastest)
+    - fastvit_t12: 7M params
+    - fastvit_s12: 9M params
+    - fastvit_sa12: 21M params (with attention)
+
+**Requirements**:
+    - timm >= 0.9.0 (for FastViT models)
+
+Reference:
+    Vasu, P.K.A., et al. (2023). FastViT: A Fast Hybrid Vision Transformer
+    using Structural Reparameterization. ICCV 2023.
+    https://arxiv.org/abs/2303.14189
+
+Author: Ductho Le (ductho.le@outlook.com)
+"""
+
+import torch
+import torch.nn as nn
+
+from wavedl.models._timm_utils import build_regression_head
+from wavedl.models.base import BaseModel
+from wavedl.models.registry import register_model
+
+
+__all__ = [
+    "FastViTBase",
+    "FastViTS12",
+    "FastViTSA12",
+    "FastViTT8",
+    "FastViTT12",
+]
+
+
+# =============================================================================
+# FASTVIT BASE CLASS
+# =============================================================================
+
+
+class FastViTBase(BaseModel):
+    """
+    FastViT base class wrapping timm implementation.
+
+    Uses RepMixer for efficient token mixing with reparameterization.
+    2D only.
+    """
+
+    def __init__(
+        self,
+        in_shape: tuple[int, int],
+        out_size: int,
+        model_name: str = "fastvit_t8",
+        pretrained: bool = True,
+        freeze_backbone: bool = False,
+        dropout_rate: float = 0.3,
+        **kwargs,
+    ):
+        super().__init__(in_shape, out_size)
+
+        if len(in_shape) != 2:
+            raise ValueError(f"FastViT requires 2D input (H, W), got {len(in_shape)}D")
+
+        self.pretrained = pretrained
+        self.freeze_backbone = freeze_backbone
+        self.model_name = model_name
+
+        # Try to load from timm
+        try:
+            import timm
+
+            self.backbone = timm.create_model(
+                model_name,
+                pretrained=pretrained,
+                num_classes=0,  # Remove classifier
+            )
+
+            # Get feature dimension
+            with torch.no_grad():
+                dummy = torch.zeros(1, 3, *in_shape)
+                features = self.backbone(dummy)
+                in_features = features.shape[-1]
+
+        except ImportError:
+            raise ImportError(
+                "timm >= 0.9.0 is required for FastViT. "
+                "Install with: pip install timm>=0.9.0"
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failed to load FastViT model '{model_name}': {e}")
+
+        # Adapt input channels (3 -> 1)
+        self._adapt_input_channels()
+
+        # Regression head
+        self.head = build_regression_head(in_features, out_size, dropout_rate)
+
+        if freeze_backbone:
+            self._freeze_backbone()
+
+    def _adapt_input_channels(self):
+        """Adapt all conv layers with 3 input channels for single-channel input."""
+        # FastViT may have multiple modules with 3 input channels (e.g., conv_kxk, conv_scale)
+        # We need to adapt all of them
+        adapted_count = 0
+
+        for name, module in self.backbone.named_modules():
+            if hasattr(module, "in_channels") and module.in_channels == 3:
+                # Check if this is a wrapper (e.g., ConvNormAct) with inner .conv
+                if hasattr(module, "conv") and isinstance(module.conv, nn.Conv2d):
+                    # Adapt the inner conv layer
+                    old_conv = module.conv
+                    module.conv = self._make_new_conv(old_conv)
+                    adapted_count += 1
+                elif isinstance(module, nn.Conv2d):
+                    # Direct Conv2d - replace it
+                    parts = name.split(".")
+                    parent = self.backbone
+                    for part in parts[:-1]:
+                        parent = getattr(parent, part)
+                    child_name = parts[-1]
+                    new_conv = self._make_new_conv(module)
+                    setattr(parent, child_name, new_conv)
+                    adapted_count += 1
+
+        if adapted_count == 0:
+            import warnings
+
+            warnings.warn(
+                "Could not adapt FastViT input channels. Model may fail.", stacklevel=2
+            )
+
+    def _make_new_conv(self, old_conv: nn.Conv2d) -> nn.Conv2d:
+        """Create new conv layer with 1 input channel."""
+        new_conv = nn.Conv2d(
+            1,
+            old_conv.out_channels,
+            kernel_size=old_conv.kernel_size,
+            stride=old_conv.stride,
+            padding=old_conv.padding,
+            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)
+        return new_conv
+
+    def _freeze_backbone(self):
+        """Freeze backbone parameters."""
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        features = self.backbone(x)
+        return self.head(features)
+
+    def reparameterize(self):
+        """
+        Reparameterize model for inference.
+
+        Fuses RepMixer blocks for faster inference.
+        Call this before deployment.
+        """
+        if hasattr(self.backbone, "reparameterize"):
+            self.backbone.reparameterize()
+
+
+# =============================================================================
+# REGISTERED VARIANTS
+# =============================================================================
+
+
+@register_model("fastvit_t8")
+class FastViTT8(FastViTBase):
+    """
+    FastViT-T8: ~3.3M backbone parameters (fastest variant).
+
+    Optimized for mobile and edge deployment.
+    2D only.
+
+    Example:
+        >>> model = FastViTT8(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_name="fastvit_t8",
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"FastViT_T8(in_shape={self.in_shape}, out_size={self.out_size}, "
+            f"pretrained={self.pretrained})"
+        )
+
+
+@register_model("fastvit_t12")
+class FastViTT12(FastViTBase):
+    """
+    FastViT-T12: ~6.5M backbone parameters.
+
+    Balanced speed and accuracy.
+    2D only.
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_name="fastvit_t12",
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"FastViT_T12(in_shape={self.in_shape}, out_size={self.out_size}, "
+            f"pretrained={self.pretrained})"
+        )
+
+
+@register_model("fastvit_s12")
+class FastViTS12(FastViTBase):
+    """
+    FastViT-S12: ~8.5M backbone parameters.
+
+    Slightly larger for better accuracy.
+    2D only.
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_name="fastvit_s12",
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"FastViT_S12(in_shape={self.in_shape}, out_size={self.out_size}, "
+            f"pretrained={self.pretrained})"
+        )
+
+
+@register_model("fastvit_sa12")
+class FastViTSA12(FastViTBase):
+    """
+    FastViT-SA12: ~10.6M backbone parameters.
+
+    With self-attention for better accuracy at the cost of speed.
+    2D only.
+    """
+
+    def __init__(self, in_shape: tuple[int, int], out_size: int, **kwargs):
+        super().__init__(
+            in_shape=in_shape,
+            out_size=out_size,
+            model_name="fastvit_sa12",
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"FastViT_SA12(in_shape={self.in_shape}, out_size={self.out_size}, "
+            f"pretrained={self.pretrained})"
+        )