| """Fake-Mamba network (XLS-R 300M + PN-BiMamba encoder), ported from |
| https://github.com/xuanxixi/Fake-Mamba (no LICENSE upstream) @ e966d2f: |
| |
| - model.py:SSLModel -> SSLModel |
| - model.py:Fake_Mamba -> FakeMambaNet |
| - conformer00.py:PN_BiMambas_Encoder -> PN_BiMambas_Encoder |
| - conformer00.py:PN_BiMambas -> PN_BiMambas |
| - mamba_ssm 1.x modules/mamba_simple.Mamba + ops/selective_scan_interface |
| .selective_scan_ref (vendored subset, as in benchmarks/RawBMamba) -> Mamba, selective_scan_ref |
| |
| Parameter / module names are identical to the source so the released Exp65 |
| ``best.pth`` loads with strict=True. |
| |
| Edits vs. the source (none changes the math): |
| 1. model.py imports two modules that are not in the repo |
| (``conformer.ConformerBlock, BiMambaEncoder``; conformer00 imports |
| ``ConExBiMamba.ExBiMamba``). They are used only by other back-ends, so only |
| Fake_Mamba + PN_BiMambas are ported; nothing is stubbed. |
| 2. ``PN_BiMambas(dim=144, depth=12)``: the released checkpoint Exp65 has 12 |
| PN-BiMamba blocks (encoder.layers.0-11); the source/paper default is 7. |
| 3. SSLModel builds the wav2vec2 XLS-R 300M architecture from its stored model |
| config (``xlsr2_300m_cfg.json``, extracted from fairseq's xlsr2_300m.pt and |
| merged onto the installed Wav2Vec2Config -- the cross-version fix used in |
| W2V2-AASIST) instead of ``load_model_ensemble_and_task`` on a hard-coded |
| path; the fine-tuned checkpoint overwrites every weight. ``extract_feat`` |
| no longer calls ``self.model.train()`` on a device/dtype mismatch (it would |
| re-enable dropout at inference); the wrapper moves the net to the device |
| before ``.eval()`` anyway. |
| 4. Mamba: mamba_ssm's CUDA fast path (``mamba_inner_fn`` / causal_conv1d / |
| selective_scan_cuda) is replaced by the library's own pure-PyTorch |
| reference (``use_fast_path=False`` branch: nn.Conv1d + ``selective_scan_ref``), |
| identical math; runs on every GPU (incl. sm_120) and exports to ONNX/TensorRT. |
| The inference-cache (step) path is dropped (unused). The time recurrence |
| is evaluated with a Hillis-Steele parallel prefix scan |
| (``selective_scan_parallel``; same recurrence, float-rounding-level |
| differences, ~log2(L) tensor ops instead of an L-step loop); |
| ``FAKE_MAMBA_SCAN=ref`` selects the library's sequential loop (parity is |
| checked in the tests). |
| 5. PN_BiMambas_Encoder runs the forward and the flipped (backward) sequence |
| through the shared Mamba in ONE call (concatenated along the batch axis) |
| instead of two calls -- same weights, same per-sample math, half the |
| Python scan loop. |
| 6. ``print`` calls removed; ``return_embedding`` kept. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import math |
| import os |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from einops import rearrange, repeat |
|
|
| _DIR = os.path.dirname(os.path.abspath(__file__)) |
| _XLSR_CFG = os.path.join(_DIR, "xlsr2_300m_cfg.json") |
|
|
| WINDOW = 66800 |
|
|
|
|
| def pad_repeat(x: np.ndarray, max_len: int = WINDOW) -> np.ndarray: |
| """source utils.pad: first max_len samples; shorter clips tiled (repeat).""" |
| x = np.asarray(x, dtype=np.float32).reshape(-1) |
| n = x.shape[0] |
| if n >= max_len: |
| return x[:max_len] |
| reps = int(max_len / n) + 1 |
| return np.tile(x, reps)[:max_len].astype(np.float32) |
|
|
|
|
| |
| class SSLModel(nn.Module): |
| def __init__(self, cfg_path: str = _XLSR_CFG): |
| super().__init__() |
| from fairseq.dataclass.utils import merge_with_parent |
| from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Config, Wav2Vec2Model |
| from omegaconf import OmegaConf |
|
|
| with open(cfg_path) as f: |
| stored = OmegaConf.create(json.load(f)) |
| cfg = merge_with_parent(Wav2Vec2Config(), stored, remove_missing=True) |
| self.model = Wav2Vec2Model.build_model(cfg, task=None) |
| self.out_dim = 1024 |
|
|
| def extract_feat(self, input_data): |
| |
| input_tmp = input_data[:, :, 0] if input_data.ndim == 3 else input_data |
| return self.model(input_tmp, mask=False, features_only=True)["x"] |
|
|
|
|
| |
| def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False): |
| """mamba_ssm 1.x ops/selective_scan_interface.selective_scan_ref (real A, variable B/C). |
| |
| u, delta, z: (b, d, l); A: (d, n); B, C: (b, n, l); D, delta_bias: (d,) |
| """ |
| dtype_in = u.dtype |
| u = u.float() |
| delta = delta.float() |
| if delta_bias is not None: |
| delta = delta + delta_bias[..., None].float() |
| if delta_softplus: |
| delta = F.softplus(delta) |
| batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1] |
| B = B.float() |
| C = C.float() |
| x = A.new_zeros((batch, dim, dstate)) |
| ys = [] |
| deltaA = torch.exp(torch.einsum("bdl,dn->bdln", delta, A)) |
| deltaB_u = torch.einsum("bdl,bnl,bdl->bdln", delta, B, u) |
| for i in range(u.shape[2]): |
| x = deltaA[:, :, i] * x + deltaB_u[:, :, i] |
| y = torch.einsum("bdn,bn->bd", x, C[:, :, i]) |
| ys.append(y) |
| y = torch.stack(ys, dim=2) |
| out = y if D is None else y + u * rearrange(D, "d -> d 1") |
| if z is not None: |
| out = out * F.silu(z) |
| return out.to(dtype=dtype_in) |
|
|
|
|
| def selective_scan_parallel(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False): |
| """Same recurrence as ``selective_scan_ref`` (h_t = exp(dt_t A) h_{t-1} + dt_t B_t u_t, |
| y_t = C_t . h_t), evaluated as a chunked parallel scan: a Hillis-Steele prefix scan |
| inside chunks of _SCAN_CHUNK steps (log2 levels of pure tensor ops), then a short |
| sequential carry across chunks -- instead of an L-step Python loop. The combine |
| (a1, b1) o (a2, b2) = (a1 a2, a2 b1 + b2) only multiplies decays <= 1, so it is |
| numerically stable; results equal the sequential scan up to float rounding. |
| """ |
| dtype_in = u.dtype |
| u = u.float() |
| delta = delta.float() |
| if delta_bias is not None: |
| delta = delta + delta_bias[..., None].float() |
| if delta_softplus: |
| delta = F.softplus(delta) |
| L = u.shape[2] |
| K = _SCAN_CHUNK |
| nc = -(-L // K) |
| |
| a = torch.exp(torch.einsum("bdl,dn->lbdn", delta, A)) |
| h = torch.einsum("bdl,bnl,bdl->lbdn", delta, B.float(), u) |
| if nc * K != L: |
| a = torch.cat([a, a.new_ones((nc * K - L,) + a.shape[1:])]) |
| h = torch.cat([h, h.new_zeros((nc * K - L,) + h.shape[1:])]) |
| a = a.reshape((nc, K) + a.shape[1:]) |
| h = h.reshape((nc, K) + h.shape[1:]) |
| |
| k = 1 |
| while k < K: |
| |
| h = torch.cat([h[:, :k], h[:, k:] + a[:, k:] * h[:, :-k]], dim=1) |
| a = torch.cat([a[:, :k], a[:, k:] * a[:, :-k]], dim=1) |
| k *= 2 |
| |
| carry = [h.new_zeros(h.shape[2:])] |
| for c in range(nc - 1): |
| carry.append(h[c, -1] + a[c, -1] * carry[-1]) |
| h = h + a * torch.stack(carry)[:, None] |
| h = h.reshape((nc * K,) + h.shape[2:])[:L] |
| y = torch.einsum("lbdn,bnl->bdl", h, C.float()) |
| out = y if D is None else y + u * rearrange(D, "d -> d 1") |
| if z is not None: |
| out = out * F.silu(z) |
| return out.to(dtype=dtype_in) |
|
|
|
|
| _SCAN_CHUNK = 8 |
|
|
| |
| SCAN_IMPL = os.environ.get("FAKE_MAMBA_SCAN", "parallel") |
|
|
|
|
| class Mamba(nn.Module): |
| """mamba_ssm 1.x modules/mamba_simple.Mamba (forward, use_fast_path=False branch).""" |
|
|
| def __init__(self, d_model, d_state=16, d_conv=4, expand=2, dt_rank="auto", |
| dt_min=0.001, dt_max=0.1, dt_init="random", dt_scale=1.0, dt_init_floor=1e-4, |
| conv_bias=True, bias=False): |
| super().__init__() |
| self.d_model = d_model |
| self.d_state = d_state |
| self.d_conv = d_conv |
| self.expand = expand |
| self.d_inner = int(self.expand * self.d_model) |
| self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank |
|
|
| self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=bias) |
| self.conv1d = nn.Conv1d(self.d_inner, self.d_inner, bias=conv_bias, kernel_size=d_conv, |
| groups=self.d_inner, padding=d_conv - 1) |
| self.activation = "silu" |
| self.act = nn.SiLU() |
| self.x_proj = nn.Linear(self.d_inner, self.dt_rank + self.d_state * 2, bias=False) |
| self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True) |
| |
| A = repeat(torch.arange(1, self.d_state + 1, dtype=torch.float32), "n -> d n", d=self.d_inner) |
| self.A_log = nn.Parameter(torch.log(A).contiguous()) |
| self.D = nn.Parameter(torch.ones(self.d_inner)) |
| self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias) |
|
|
| def forward(self, hidden_states): |
| batch, seqlen, dim = hidden_states.shape |
| xz = rearrange(self.in_proj.weight @ rearrange(hidden_states, "b l d -> d (b l)"), |
| "d (b l) -> b d l", l=seqlen) |
| if self.in_proj.bias is not None: |
| xz = xz + rearrange(self.in_proj.bias.to(dtype=xz.dtype), "d -> d 1") |
| A = -torch.exp(self.A_log.float()) |
| x, z = xz.chunk(2, dim=1) |
| x = self.act(self.conv1d(x)[..., :seqlen]) |
| x_dbl = self.x_proj(rearrange(x, "b d l -> (b l) d")) |
| dt, B, C = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1) |
| dt = self.dt_proj.weight @ dt.t() |
| dt = rearrange(dt, "d (b l) -> b d l", l=seqlen) |
| B = rearrange(B, "(b l) dstate -> b dstate l", l=seqlen).contiguous() |
| C = rearrange(C, "(b l) dstate -> b dstate l", l=seqlen).contiguous() |
| scan = selective_scan_parallel if SCAN_IMPL == "parallel" else selective_scan_ref |
| y = scan(x, dt, A, B, C, self.D.float(), z=z, |
| delta_bias=self.dt_proj.bias.float(), delta_softplus=True) |
| y = rearrange(y, "b d l -> b l d") |
| return self.out_proj(y) |
|
|
|
|
| |
| class PN_BiMambas_Encoder(nn.Module): |
| def __init__(self, d_model, n_state): |
| super().__init__() |
| self.d_model = d_model |
| self.mamba = Mamba(d_model, n_state) |
| self.norm1 = nn.LayerNorm(d_model) |
| self.norm2 = nn.LayerNorm(d_model) |
| self.feed_forward = nn.Sequential( |
| nn.Linear(d_model, d_model * 4), |
| nn.GELU(), |
| nn.Linear(d_model * 4, d_model), |
| ) |
|
|
| def forward(self, x): |
| residual = x |
| x_norm = self.norm1(x) |
| b = x_norm.shape[0] |
| |
| both = self.mamba(torch.cat([x_norm, torch.flip(x_norm, dims=[1])], dim=0)) |
| mamba_out_forward = both[:b] |
| mamba_out_backward = torch.flip(both[b:], dims=[1]) |
| mamba_out = mamba_out_forward + mamba_out_backward |
| mamba_out = self.norm2(mamba_out) |
| ff_out = self.feed_forward(mamba_out) |
| return ff_out + residual |
|
|
|
|
| class PN_BiMambas(nn.Module): |
| def __init__(self, dim, *, depth): |
| super().__init__() |
| self.dim = dim |
| self.layers = nn.ModuleList([PN_BiMambas_Encoder(d_model=dim, n_state=16) for _ in range(depth)]) |
|
|
| def forward(self, x): |
| for block in self.layers: |
| x = block(x) |
| return x |
|
|
|
|
| class FakeMambaNet(nn.Module): |
| """model.py:Fake_Mamba (args.emb_size = 144). forward(x[B, T]) -> logits[B, 2].""" |
|
|
| def __init__(self, emb_size: int = 144, depth: int = 12): |
| super().__init__() |
| self.ssl_model = SSLModel() |
| self.LL = nn.Linear(1024, emb_size) |
| self.first_bn = nn.BatchNorm2d(num_features=1) |
| self.selu = nn.SELU(inplace=True) |
| self.encoder = PN_BiMambas(dim=144, depth=depth) |
| self.attention_pool = nn.Linear(144, 1) |
| self.fc5 = nn.Linear(144, 2) |
|
|
| def forward(self, x, return_embedding=False): |
| x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1)) |
| x = self.LL(x_ssl_feat) |
| x = x.unsqueeze(dim=1) |
| x = self.first_bn(x) |
| x = self.selu(x) |
| x = x.squeeze(dim=1) |
| x = self.encoder(x) |
| attention_weights = F.softmax(self.attention_pool(x), dim=1) |
| x_pooled = torch.matmul(attention_weights.transpose(-1, -2), x).squeeze(-2) |
| if return_embedding: |
| return self.fc5(x_pooled), x_pooled |
| return self.fc5(x_pooled) |
|
|
|
|
| def load_state(path: str) -> dict: |
| sd = torch.load(path, map_location="cpu") |
| return sd.get("state_dict", sd) if isinstance(sd, dict) and "state_dict" in sd else sd |
|
|
