File size: 4,233 Bytes
ae9d1a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """SSL-Mamba network (inference-only port).
Source: github.com/hoanmyTran/spoofing_ssl_mamba @ 9afc33de7f74a9b9b3423a9028c10e1d18c1c6b8 (MIT),
``model.py`` ``SSLClassifier``. Checkpoint: HF hoanmyTran/ssl_spoofing_mamba @
4e72dfcef4f03460e08f68a4a5fbcf3705a3414e ``interspeech2025.pt`` (Apache-2.0; flat state_dict),
converted once to ``model.safetensors`` with the training-only ``loss_fn.weight`` dropped.
Edits vs the source (all inference-only; module names unchanged so keys load strictly):
1. pytorch_lightning / torchmetrics scaffolding, loss and training/eval loops removed.
2. XLS-R 300M built from the bundled config (``xlsr_300m_config.json`` = the
facebook/wav2vec2-xls-r-300m config) instead of ``from_pretrained``: no Hub weight
download; the checkpoint overwrites every weight (the front-end is fully fine-tuned).
3. ``mamba_ssm.Mamba`` -> ``_mamba_v1.Mamba`` (pure PyTorch, same parameter names, same
math; chunked closed-form selective scan, see that file). mamba_ssm / causal_conv1d are
not installed and Triton is unreliable on sm_120.
4. ``multiconv_cgmlp`` espnet imports -> ``_espnet_shim`` (LayerNorm eps 1e-12, activations);
``pooling.py`` (MultiHeadAttentionPooling) vendored from the authors' sibling repo
(dissimilarity_deepfake_detection, as in benchmarks/SSL-Feature-Gating/), not in this repo.
5. ``nn.SiLU(inplace=True)`` / ``nn.SELU(inplace=True)`` -> non-inplace (export-friendly;
identical values).
Forward is otherwise verbatim: 25 hidden states -> Linear(1024,128) -> SiLU -> sum over layers
-> MultiConv cgMLP (7/15/23/31) -> 12 x [x = x + LayerNorm(Mamba(x))] (unidirectional)
-> MHAP (4 heads, mean+std) -> Linear 256->128 -> SELU -> Linear 128->2.
"""
from __future__ import annotations
import os
import numpy as np
import torch
import torch.nn as nn
from transformers import Wav2Vec2Config, Wav2Vec2Model
from _mamba_v1 import Mamba
from multiconv_cgmlp import MultiConvolutionalGatingMLP
from pooling import MultiHeadAttentionPooling
_DIR = os.path.dirname(os.path.abspath(__file__))
_CFG = os.path.join(_DIR, "xlsr_300m_config.json")
WINDOW = 64600 # source config dataset.max_len (4.0375 s @ 16 kHz)
def pad_repeat(x: np.ndarray, max_len: int = WINDOW) -> np.ndarray:
"""Source dataset.pad: first max_len samples; shorter clips tile-repeated."""
x = np.asarray(x, dtype=np.float32).reshape(-1)
n = x.shape[0]
if n >= max_len:
return x[:max_len]
return np.tile(x, max_len // n + 1)[:max_len].astype(np.float32)
class SSLMambaNet(nn.Module):
def __init__(self):
super().__init__()
self.feature_extractor = Wav2Vec2Model(Wav2Vec2Config.from_json_file(_CFG))
self.cgmlp = MultiConvolutionalGatingMLP(
size=128, linear_units=1024, arch_type="concat_fusion", kernel_sizes="7,15,23,31",
merge_conv_kernel=31, use_non_linear=True, dropout_rate=0.1,
use_linear_after_conv=True, activation="silu", gate_activation="silu")
self.feature_projection = nn.Linear(1024, 128)
self.blocks = nn.ModuleList(
nn.Sequential(Mamba(d_model=128, d_state=16, d_conv=4, expand=2), nn.LayerNorm(128))
for _ in range(12))
self.last_pooling = MultiHeadAttentionPooling(128)
self.silu = nn.SiLU()
self.classifier = nn.Sequential(nn.Linear(256, 128), nn.SELU(), nn.Linear(128, 2))
def set_ref_scan(self, on: bool) -> None:
for m in self.modules():
if isinstance(m, Mamba):
m.use_ref_scan = on
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.feature_extractor(x, output_hidden_states=True)
hs = torch.stack(x.hidden_states, dim=1) # (B, 25, T, 1024)
hs = self.silu(self.feature_projection(hs))
x = torch.sum(hs, 1) # (B, T, 128)
x = self.cgmlp(x)
for b in self.blocks:
x = x + b(x)
x = self.last_pooling(x.permute(0, 2, 1)) # (B, 256, 1)
return self.classifier(x.squeeze(-1)) # (B, 2); class 1 = bona fide
def load_state(path: str) -> dict:
from safetensors.torch import load_file
return load_file(path)
|