dloc-code / mamba_model.py

Upload folder using huggingface_hub

a6b4c35 verified 6 months ago

10.8 kB

#!/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

	#!/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