Wrapper code (speech-spoof-bench)

29161c2 verified 6 days ago

14.3 kB

"""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

	"""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