Delete duplicate/old file: \/mamba_model.py

///mamba_model.py

@@ -1,273 +0,0 @@
-#!/usr/bin/env python
-"""
-mamba_model.py
-Mamba-DLoc: Replaces the ResNet bottleneck in DLoc with Mamba selective
-state-space (S6) blocks for efficient long-range sequence modelling.
-
-Architecture overview
----------------------
-CNN Stem  (4 -> 128 channels, 2x spatial downsample)
-  -> Tokenize (flatten spatial -> sequence)
-  -> 4x MambaResidualBlock  (LayerNorm -> Mamba -> residual)
-  -> Reshape back to spatial
-  -> CNN Decoder (upsample back to original resolution)
-
-Input:  [B, 4, 161, 361]  AoA heatmaps
-Output: [B, 1, 161, 361]  location heatmap
-"""
-
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-# ---------------------------------------------------------------------------
-# Try to import the CUDA-optimised Mamba implementation (Linux + CUDA only).
-# Fall back to a pure-PyTorch implementation on Windows / CPU.
-# ---------------------------------------------------------------------------
-HAS_MAMBA_SSM = False
-try:
-    from mamba_ssm import Mamba as CudaMamba  # type: ignore
-    HAS_MAMBA_SSM = True
-except ImportError:
-    CudaMamba = None
-
-
-# ===================================================================
-# Pure-PyTorch Mamba (S6 selective scan) - works on any backend
-# ===================================================================
-class PurePytorchMamba(nn.Module):
-    """
-    Minimal re-implementation of the Mamba selective-scan block (S6).
-
-    Interface: [B, L, D] -> [B, L, D]
-
-    Parameters
-    ----------
-    d_model : int   - input / output dimension (D)
-    d_state : int   - SSM hidden state dimension (N)
-    d_conv  : int   - local convolution kernel width
-    expand  : int   - expansion factor for inner dimension (E = expand * D)
-    """
-
-    def __init__(self, d_model: int = 128, d_state: int = 16,
-                 d_conv: int = 4, expand: int = 2):
-        super().__init__()
-        self.d_model = d_model
-        self.d_state = d_state
-        self.d_conv = d_conv
-        self.expand = expand
-        self.d_inner = d_model * expand           # E
-        self.dt_rank = math.ceil(d_model / 16)     # rank for delta-t projection
-
-        # --- input projection: D -> 2*E  (main path + gate) ---
-        self.in_proj = nn.Linear(d_model, 2 * self.d_inner, bias=False)
-
-        # --- causal depthwise conv on main path ---
-        self.conv1d = nn.Conv1d(
-            self.d_inner, self.d_inner,
-            kernel_size=d_conv, padding=d_conv - 1,  # causal: trim later
-            groups=self.d_inner, bias=True,
-        )
-
-        # --- SSM parameter projections ---
-        #   from the convolved main path (E) -> dt_rank + 2*N
-        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + 2 * d_state, bias=False)
-
-        #   dt_rank -> E  (broadcast delta across inner dim)
-        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
-
-        # --- learnable SSM parameters ---
-        # A_log: [E, N]  initialised as log(1..N) broadcast to each inner dim
-        A = torch.arange(1, d_state + 1, dtype=torch.float32).unsqueeze(0).expand(self.d_inner, -1)
-        self.A_log = nn.Parameter(torch.log(A))           # learned in log-space
-
-        # D (skip connection scalar per inner dim)
-        self.D = nn.Parameter(torch.ones(self.d_inner))
-
-        # --- output projection: E -> D ---
-        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
-
-    # ---------------------------------------------------------------
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        x : [B, L, D]
-        returns : [B, L, D]
-        """
-        B, L, _ = x.shape
-
-        # --- input projection & split ---
-        xz = self.in_proj(x)                              # [B, L, 2E]
-        x_main, z = xz.chunk(2, dim=-1)                   # each [B, L, E]
-
-        # --- causal conv1d on main path ---
-        # conv1d expects [B, E, L]
-        x_main = x_main.transpose(1, 2)                   # [B, E, L]
-        x_main = self.conv1d(x_main)[:, :, :L]            # causal trim
-        x_main = x_main.transpose(1, 2)                   # [B, L, E]
-        x_main = F.silu(x_main)
-
-        # --- SSM parameters ---
-        x_ssm = self.x_proj(x_main)                       # [B, L, dt_rank + 2N]
-        dt, B_param, C_param = x_ssm.split(
-            [self.dt_rank, self.d_state, self.d_state], dim=-1
-        )
-        # dt: [B, L, dt_rank] -> [B, L, E]
-        dt = F.softplus(self.dt_proj(dt))                  # [B, L, E]
-
-        # A: [E, N]  (negative for stability)
-        A = -torch.exp(self.A_log)                         # [E, N]
-
-        # --- discretize ---
-        # dA: [B, L, E, N]
-        dA = torch.exp(dt.unsqueeze(-1) * A.unsqueeze(0).unsqueeze(0))
-        # dB: [B, L, E, N]
-        dB = dt.unsqueeze(-1) * B_param.unsqueeze(2)       # broadcast E
-
-        # --- selective scan (sequential, kept simple for correctness) ---
-        # x_main: [B, L, E]  ->  weight input
-        # h: [B, E, N]       SSM hidden state
-        h = torch.zeros(B, self.d_inner, self.d_state,
-                        device=x.device, dtype=x.dtype)
-        ys = []
-        for t in range(L):
-            h = dA[:, t, :, :] * h + dB[:, t, :, :] * x_main[:, t, :].unsqueeze(-1)
-            y_t = (h * C_param[:, t, :].unsqueeze(1)).sum(dim=-1)   # [B, E]
-            ys.append(y_t)
-        y = torch.stack(ys, dim=1)                         # [B, L, E]
-
-        # --- skip connection ---
-        y = y + x_main * self.D.unsqueeze(0).unsqueeze(0)
-
-        # --- gate ---
-        y = y * F.silu(z)
-
-        # --- output projection ---
-        return self.out_proj(y)                            # [B, L, D]
-
-
-# ===================================================================
-# Mamba residual block: LayerNorm -> Mamba -> residual add
-# ===================================================================
-class MambaResidualBlock(nn.Module):
-    """
-    Pre-norm residual block wrapping a Mamba layer.
-
-    Interface: [B, L, D] -> [B, L, D]
-    Uses the CUDA Mamba kernel when available, otherwise falls back to
-    PurePytorchMamba.
-    """
-
-    def __init__(self, d_model: int = 128, d_state: int = 16,
-                 d_conv: int = 4, expand: int = 2):
-        super().__init__()
-        self.norm = nn.LayerNorm(d_model)
-        if HAS_MAMBA_SSM and CudaMamba is not None:
-            self.mamba = CudaMamba(
-                d_model=d_model, d_state=d_state,
-                d_conv=d_conv, expand=expand,
-            )
-        else:
-            self.mamba = PurePytorchMamba(
-                d_model=d_model, d_state=d_state,
-                d_conv=d_conv, expand=expand,
-            )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x + self.mamba(self.norm(x))
-
-
-# ===================================================================
-# MambaDLocNet - full end-to-end model
-# ===================================================================
-class MambaDLocNet(nn.Module):
-    """
-    Mamba-DLoc: CNN stem  ->  Mamba sequence blocks  ->  CNN decoder.
-
-    Input:  [B, 4, 161, 361]
-    Output: [B, 1, 161, 361]
-    """
-
-    # Spatial sizes at each resolution level (full, mid, low)
-    SPATIAL_SIZES = [(161, 361), (81, 181), (41, 91)]
-
-    def __init__(self, input_nc: int = 4, d_model: int = 128,
-                 n_mamba_blocks: int = 4,
-                 d_state: int = 16, d_conv: int = 4, expand: int = 2):
-        super().__init__()
-        self.d_model = d_model
-
-        # ------ CNN Stem (encoder) ------
-        self.stem = nn.Sequential(
-            # [B,input_nc,161,361] -> [B,32,161,361]
-            nn.Conv2d(input_nc, 32, kernel_size=3, stride=1, padding=1, bias=False),
-            nn.BatchNorm2d(32),
-            nn.ReLU(inplace=True),
-            # [B,32,161,361] -> [B,64,81,181]
-            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1, bias=False),
-            nn.BatchNorm2d(64),
-            nn.ReLU(inplace=True),
-            # [B,64,81,181] -> [B,128,41,91]
-            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False),
-            nn.BatchNorm2d(128),
-            nn.ReLU(inplace=True),
-        )
-
-        # ------ Mamba sequence blocks ------
-        self.mamba_blocks = nn.Sequential(*[
-            MambaResidualBlock(d_model=d_model, d_state=d_state,
-                               d_conv=d_conv, expand=expand)
-            for _ in range(n_mamba_blocks)
-        ])
-
-        # ------ CNN Decoder ------
-        # Stage 1: [B,128,41,91] -> interpolate to (81,181) -> conv 128->64
-        self.dec_conv1 = nn.Sequential(
-            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1, bias=False),
-            nn.BatchNorm2d(64),
-            nn.ReLU(inplace=True),
-        )
-        # Stage 2: [B,64,81,181] -> interpolate to (161,361) -> conv 64->32
-        self.dec_conv2 = nn.Sequential(
-            nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1, bias=False),
-            nn.BatchNorm2d(32),
-            nn.ReLU(inplace=True),
-        )
-        # Stage 3: 1x1 head -> sigmoid
-        self.dec_head = nn.Sequential(
-            nn.Conv2d(32, 1, kernel_size=1, bias=True),
-            nn.Sigmoid(),
-        )
-
-    # ---------------------------------------------------------------
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        x : [B, 4, 161, 361]
-        returns : [B, 1, 161, 361]
-        """
-        B = x.size(0)
-
-        # --- CNN Stem ---
-        feat = self.stem(x)                                # [B, 128, 41, 91]
-
-        # --- Tokenize: flatten spatial dims into sequence ---
-        _, C, H, W = feat.shape
-        tokens = feat.reshape(B, C, H * W).permute(0, 2, 1)  # [B, 3731, 128]
-
-        # --- Mamba blocks ---
-        tokens = self.mamba_blocks(tokens)                 # [B, 3731, 128]
-
-        # --- Reshape back to spatial ---
-        feat = tokens.permute(0, 2, 1).reshape(B, C, H, W)   # [B, 128, 41, 91]
-
-        # --- CNN Decoder ---
-        feat = F.interpolate(feat, size=self.SPATIAL_SIZES[1],
-                             mode='bilinear', align_corners=False)   # [B,128,81,181]
-        feat = self.dec_conv1(feat)                                   # [B,64,81,181]
-
-        feat = F.interpolate(feat, size=self.SPATIAL_SIZES[0],
-                             mode='bilinear', align_corners=False)   # [B,64,161,361]
-        feat = self.dec_conv2(feat)                                   # [B,32,161,361]
-
-        out = self.dec_head(feat)                                     # [B,1,161,361]
-        return out