"""XLSR-Nes2Net-LA (Nes2Net-X + local-window cross-block attention) for the benchmark harness. Source: github XuepingZhang/MultiAPI-Spoof @ 1f712e3fff26e237268385b92365a9debd4cb69f, models/XLSR2_Nes2Net_X_ATT.py (CC BY-NC 4.0). Config job/xlsr2_nes2net_ATT/few_dkucc.conf: Nes_ratio [8,8], SE_ratio [1], dilation 2, pool_func mean, n_output_logits 2. Edits vs the source (math unchanged): - SSLModel: built from the local xlsr2_300m.pt via the fairseq cross-version fix (benchmarks/W2V2-AASIST/_net.py, reused through benchmarks/Nes2Net/_net.py) instead of load_model_ensemble_and_task with a hard-coded cluster path. - Nested_Res2Net_TDNN still constructs the source's `global_attn` (nn.MultiheadAttention(1024, 8)), which forward() never calls; it exists only so the checkpoint loads strictly. - Model takes no args dict (config above is baked in); forward returns (ssl_feat, logits) as the source. """ from __future__ import annotations import os import math import torch import torch.nn as nn import fairseq # Base XLS-R 300M checkpoint, shared with the W2V2-AASIST benchmark (symlinked). # Only used to *construct* the wav2vec2 architecture; the fine-tuned Nes2Net # checkpoint overwrites every weight, including this SSL submodule. _XLSR_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "xlsr2_300m.pt") # ---------------- wav2vec 2.0 XLS-R front-end ---------------- # Vendored from benchmarks/W2V2-AASIST/_net.py::SSLModel. We bypass # load_model_ensemble_and_task (chokes on cfg keys this fairseq version doesn't # know, e.g. `multiple_train_files`) and merge the stored model cfg onto the # installed Wav2Vec2Config defaults so cross-version field gaps are filled. class SSLModel(nn.Module): def __init__(self, device, cp_path=_XLSR_PATH): super().__init__() from fairseq.models.wav2vec.wav2vec2 import Wav2Vec2Model, Wav2Vec2Config from fairseq.dataclass.utils import merge_with_parent state = fairseq.checkpoint_utils.load_checkpoint_to_cpu(cp_path) merged_cfg = merge_with_parent(Wav2Vec2Config(), state["cfg"]["model"], remove_missing=True) ssl = Wav2Vec2Model.build_model(merged_cfg, task=None) ssl.load_state_dict(state["model"], strict=True) self.model = ssl self.device = device self.out_dim = 1024 def extract_feat(self, input_data): if next(self.model.parameters()).device != input_data.device \ or next(self.model.parameters()).dtype != input_data.dtype: self.model.to(input_data.device, dtype=input_data.dtype) self.model.train() if input_data.ndim == 3: input_tmp = input_data[:, :, 0] else: input_tmp = input_data emb = self.model(input_tmp, mask=False, features_only=True)['x'] return emb # ---------------- Nes2Net-X back-end (ported verbatim) ---------------- class SEModule(nn.Module): def __init__(self, channels, SE_ratio=8): super().__init__() self.se = nn.Sequential( nn.AdaptiveAvgPool1d(1), nn.Conv1d(channels, channels // SE_ratio, kernel_size=1, padding=0), nn.ReLU(), nn.Conv1d(channels // SE_ratio, channels, kernel_size=1, padding=0), nn.Sigmoid(), ) def forward(self, input): x = self.se(input) return input * x class Bottle2neck(nn.Module): def __init__(self, inplanes, planes, kernel_size=None, dilation=None, scale=8, SE_ratio=8): super().__init__() width = int(math.floor(planes / scale)) self.conv1 = nn.Conv1d(inplanes, width * scale, kernel_size=1) self.bn1 = nn.BatchNorm1d(width * scale) self.nums = scale - 1 convs = [] bns = [] weighted_sum = [] num_pad = math.floor(kernel_size / 2) * dilation for i in range(self.nums): convs.append(nn.Conv2d(width, width, kernel_size=(kernel_size, 1), dilation=(dilation, 1), padding=(num_pad, 0))) bns.append(nn.BatchNorm2d(width)) initial_value = torch.ones(1, 1, 1, i + 2) * (1 / (i + 2)) weighted_sum.append(nn.Parameter(initial_value, requires_grad=True)) self.weighted_sum = nn.ParameterList(weighted_sum) self.convs = nn.ModuleList(convs) self.bns = nn.ModuleList(bns) self.conv3 = nn.Conv1d(width * scale, planes, kernel_size=1) self.bn3 = nn.BatchNorm1d(planes) self.relu = nn.ReLU() self.width = width self.se = SEModule(planes, SE_ratio) self.dropout = nn.Dropout(p=0.2) # source; no-op at eval def forward(self, x): residual = x out = self.conv1(x) out = self.relu(out) out = self.bn1(out).unsqueeze(-1) # bz c T 1 spx = torch.split(out, self.width, 1) sp = spx[self.nums] for i in range(self.nums): sp = torch.cat((sp, spx[i]), -1) sp = self.bns[i](self.relu(self.convs[i](sp))) sp_s = sp * self.weighted_sum[i] sp_s = torch.sum(sp_s, dim=-1, keepdim=False) if i == 0: out = sp_s else: out = torch.cat((out, sp_s), 1) out = torch.cat((out, spx[self.nums].squeeze(-1)), 1) out = self.conv3(out) out = self.relu(out) out = self.bn3(out) out = self.se(out) out = self.dropout(out) out += residual return out class ASTP(nn.Module): """Attentive statistics pooling (unused with pool_func='mean', kept for fidelity).""" def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False): super().__init__() self.global_context_att = global_context_att if global_context_att: self.linear1 = nn.Conv1d(in_dim * 3, bottleneck_dim, kernel_size=1) else: self.linear1 = nn.Conv1d(in_dim, bottleneck_dim, kernel_size=1) self.linear2 = nn.Conv1d(bottleneck_dim, in_dim, kernel_size=1) def forward(self, x): if len(x.shape) == 4: x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3]) assert len(x.shape) == 3 if self.global_context_att: context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x) context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x) x_in = torch.cat((x, context_mean, context_std), dim=1) else: x_in = x alpha = torch.tanh(self.linear1(x_in)) alpha = torch.softmax(self.linear2(alpha), dim=2) mean = torch.sum(alpha * x, dim=2) var = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2 std = torch.sqrt(var.clamp(min=1e-10)) return torch.cat([mean, std], dim=1) # ---------------- local-window cross-block attention (ported verbatim) ---------------- class LocalWindowAttention(nn.Module): """ Local window (sliding) cross-block attention. Input: xb [B, NB, C, T] Output: xb' [B, NB, C, T] (residual fused) Params: - window_size: odd integer W (number of blocks to attend for each center) - dilation: step between neighbor blocks (default 1) - num_heads: multi-head attention heads (internal MHA implemented via linear projections) - reduce_ratio: linear proj reduce factor (embed_dim = max(C//reduce_ratio, 1)) Boundary handling: - uses clamped indices (edge repeats). Could be changed to zero-pad if desired. """ def __init__(self, num_blocks, block_channels, window_size=3, dilation=1, num_heads=4, reduce_ratio=1, dropout=0.0, use_pos=True): super().__init__() assert window_size % 2 == 1, "window_size must be odd" self.NB = num_blocks self.C = block_channels self.W = window_size self.dilation = dilation self.num_heads = num_heads self.use_pos = use_pos D = max(block_channels // reduce_ratio, 1) # ensure D divisible by num_heads for simple multi-head split head_dim = max(D // num_heads, 1) D = head_dim * num_heads self.D = D self.head_dim = head_dim self.proj_q = nn.Linear(block_channels, D, bias=False) self.proj_k = nn.Linear(block_channels, D, bias=False) self.proj_v = nn.Linear(block_channels, D, bias=False) self.proj_out = nn.Linear(D, block_channels, bias=False) self.dropout = nn.Dropout(dropout) if use_pos: # position embedding per block (on D dims) self.pos_emb = nn.Parameter(torch.zeros(self.NB, D)) nn.init.trunc_normal_(self.pos_emb, std=0.02) else: self.pos_emb = None def _build_neighbor_index(self, device): """ Precompute neighbor index matrix of shape [NB, W] where each row i contains the indices (clamped) of the window centered at i with given dilation. """ half = (self.W - 1) // 2 idx = torch.arange(self.NB, device=device) # [NB] # neighbors offsets: [-half..half] * dilation offsets = torch.arange(-half, half+1, device=device) * self.dilation # [W] neighbors = idx.unsqueeze(1) + offsets.unsqueeze(0) # [NB, W] neighbors = torch.clamp(neighbors, 0, self.NB - 1) # clamp to edges (repeat edges) return neighbors.long() # [NB, W] def forward(self, xb): # xb: [B, NB, C, T] B, NB, C, T = xb.shape assert NB == self.NB and C == self.C block_tokens = xb.mean(dim=-1) # [B, NB, C] Q = self.proj_q(block_tokens) # [B, NB, D] K = self.proj_k(block_tokens) # [B, NB, D] V = self.proj_v(block_tokens) # [B, NB, D] if self.pos_emb is not None: pos = self.pos_emb.unsqueeze(0) # [1, NB, D] Q = Q + pos K = K + pos V = V + pos neighbors_idx = self._build_neighbor_index(device=xb.device) # [NB, W] K_neighbors = K.unsqueeze(1).expand(B, NB, NB, self.D).gather( 2, neighbors_idx.unsqueeze(0).unsqueeze(-1).expand(B, NB, self.W, self.D) ) V_neighbors = V.unsqueeze(1).expand(B, NB, NB, self.D).gather( 2, neighbors_idx.unsqueeze(0).unsqueeze(-1).expand(B, NB, self.W, self.D) ) Qc = Q.unsqueeze(2) # [B, NB, 1, D] scores = torch.matmul(Qc, K_neighbors.transpose(-2, -1)).squeeze(2) # [B, NB, W] scores = scores / math.sqrt(self.head_dim * self.num_heads) # scale attn = torch.softmax(scores, dim=-1) # [B, NB, W] attn = self.dropout(attn) attn_exp = attn.unsqueeze(2) context = torch.matmul(attn_exp, V_neighbors).squeeze(2) # [B, NB, D] # 6) project out and fuse out_tokens = self.proj_out(context) # [B, NB, C] # expand to time and residual add out_tokens_time = out_tokens.unsqueeze(-1).expand(-1, -1, -1, T) # [B, NB, C, T] out = xb + out_tokens_time # residual fusion return out class Nested_Res2Net_TDNN(nn.Module): def __init__(self, Nes_ratio=(8, 8), input_channel=1024, n_output_logits=2, dilation=2, pool_func='mean', SE_ratio=(1,)): super().__init__() self.Nes_ratio = Nes_ratio[0] assert input_channel % Nes_ratio[0] == 0 C = input_channel // Nes_ratio[0] self.C = C Build_in_Res2Nets = [] bns = [] for i in range(Nes_ratio[0] - 1): Build_in_Res2Nets.append( Bottle2neck(C, C, kernel_size=3, dilation=dilation, scale=Nes_ratio[1], SE_ratio=SE_ratio[0])) bns.append(nn.BatchNorm1d(C)) self.Build_in_Res2Nets = nn.ModuleList(Build_in_Res2Nets) self.bns = nn.ModuleList(bns) self.bn = nn.BatchNorm1d(1024) self.relu = nn.ReLU() self.pool_func = pool_func if pool_func == 'mean': self.fc = nn.Linear(1024, n_output_logits) elif pool_func == 'ASTP': self.pooling = ASTP(in_dim=input_channel, bottleneck_dim=128, global_context_att=False) self.fc = nn.Linear(2048, n_output_logits) self.local_block_attn = LocalWindowAttention(num_blocks=self.Nes_ratio, block_channels=self.C, window_size=3, dilation=1, num_heads=1, reduce_ratio=1, dropout=0.1, use_pos=True) # Present in the source and the checkpoint but never used by forward(). self.global_attn = nn.MultiheadAttention(embed_dim=input_channel, num_heads=8, batch_first=True) def forward(self, x): spx = torch.split(x, self.C, 1) # tuple of [B, C, T] sp_list = [] for i in range(self.Nes_ratio - 1): if i == 0: sp = spx[i] else: sp = sp + spx[i] sp = self.Build_in_Res2Nets[i](sp) sp = self.relu(sp) sp = self.bns[i](sp) sp_list.append(sp) sp_list.append(spx[-1]) # last block unchanged xb = torch.stack(sp_list, dim=1) # [B, NB, C, T] xb = self.local_block_attn(xb) out = torch.cat([xb[:, j] for j in range(self.Nes_ratio)], dim=1) out = self.bn(out) out = self.relu(out) if self.pool_func == 'mean': out = torch.mean(out, dim=-1) elif self.pool_func == 'ASTP': out = self.pooling(out) out = self.fc(out) return out class Model(nn.Module): """XLSR2_Nes2Net_X_ATT with the few_dkucc.conf model_config baked in.""" def __init__(self, device="cpu"): super().__init__() self.device = device self.ssl_model = SSLModel(device) self.Nested_Res2Net_TDNN = Nested_Res2Net_TDNN( Nes_ratio=[8, 8], input_channel=1024, n_output_logits=2, dilation=2, pool_func='mean', SE_ratio=[1]) def forward(self, x): x_ssl_feat = self.ssl_model.extract_feat(x.squeeze(-1)) x_ssl_feat = x_ssl_feat.permute(0, 2, 1) output = self.Nested_Res2Net_TDNN(x_ssl_feat) return x_ssl_feat, output