| """WavLM-large + AASIST speech deepfake detector (inference only).""" |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import WavLMModel |
|
|
|
|
| class _GraphAttentionLayer(nn.Module): |
| def __init__(self, in_dim, out_dim, temperature): |
| super().__init__() |
| self.att_proj = nn.Linear(in_dim, out_dim) |
| self.att_weight = nn.Parameter(torch.empty(out_dim, 1)) |
| self.proj_with_att = nn.Linear(in_dim, out_dim) |
| self.proj_without_att = nn.Linear(in_dim, out_dim) |
| self.bn = nn.BatchNorm1d(out_dim) |
| self.temp = temperature |
|
|
| def forward(self, x): |
| N = x.size(1) |
| x_i = x.unsqueeze(2).expand(-1, -1, N, -1) |
| att = torch.matmul(torch.tanh(self.att_proj(x_i * x_i.transpose(1, 2))), self.att_weight) / self.temp |
| att = F.softmax(att, dim=-2) |
| out = self.proj_with_att(torch.matmul(att.squeeze(-1), x)) + self.proj_without_att(x) |
| s = out.shape |
| return F.selu(self.bn(out.reshape(-1, s[-1])).reshape(s)) |
|
|
|
|
| class _HtrgGraphAttentionLayer(nn.Module): |
| def __init__(self, in_dim, out_dim, temperature): |
| super().__init__() |
| self.proj_type1 = nn.Linear(in_dim, in_dim) |
| self.proj_type2 = nn.Linear(in_dim, in_dim) |
| self.att_proj = nn.Linear(in_dim, out_dim) |
| self.att_projM = nn.Linear(in_dim, out_dim) |
| self.att_weight11 = nn.Parameter(torch.empty(out_dim, 1)) |
| self.att_weight22 = nn.Parameter(torch.empty(out_dim, 1)) |
| self.att_weight12 = nn.Parameter(torch.empty(out_dim, 1)) |
| self.att_weightM = nn.Parameter(torch.empty(out_dim, 1)) |
| self.proj_with_att = nn.Linear(in_dim, out_dim) |
| self.proj_without_att = nn.Linear(in_dim, out_dim) |
| self.proj_with_attM = nn.Linear(in_dim, out_dim) |
| self.proj_without_attM = nn.Linear(in_dim, out_dim) |
| self.bn = nn.BatchNorm1d(out_dim) |
| self.temp = temperature |
|
|
| def forward(self, x1, x2, master): |
| n1 = x1.size(1) |
| x = torch.cat([self.proj_type1(x1), self.proj_type2(x2)], dim=1) |
| N = x.size(1) |
| x_i = x.unsqueeze(2).expand(-1, -1, N, -1) |
| att_map = torch.tanh(self.att_proj(x_i * x_i.transpose(1, 2))) |
| board = torch.zeros_like(att_map[:, :, :, 0]).unsqueeze(-1) |
| board[:, :n1, :n1] = torch.matmul(att_map[:, :n1, :n1], self.att_weight11) |
| board[:, n1:, n1:] = torch.matmul(att_map[:, n1:, n1:], self.att_weight22) |
| board[:, :n1, n1:] = torch.matmul(att_map[:, :n1, n1:], self.att_weight12) |
| board[:, n1:, :n1] = torch.matmul(att_map[:, n1:, :n1], self.att_weight12) |
| board = F.softmax(board / self.temp, dim=-2) |
| out = self.proj_with_att(torch.matmul(board.squeeze(-1), x)) + self.proj_without_att(x) |
| s = out.shape |
| out = F.selu(self.bn(out.reshape(-1, s[-1])).reshape(s)) |
| att_m = F.softmax(torch.matmul(torch.tanh(self.att_projM(x * master)), self.att_weightM) / self.temp, dim=-2) |
| master = self.proj_with_attM(torch.matmul(att_m.squeeze(-1).unsqueeze(1), x)) + self.proj_without_attM(master) |
| return out.narrow(1, 0, n1), out.narrow(1, n1, N - n1), master |
|
|
|
|
| class _GraphPool(nn.Module): |
| def __init__(self, k, in_dim): |
| super().__init__() |
| self.k = k |
| self.proj = nn.Linear(in_dim, 1) |
|
|
| def forward(self, h): |
| scores = torch.sigmoid(self.proj(h)) |
| n_keep = max(int(h.size(1) * self.k), 1) |
| idx = torch.topk(scores, n_keep, dim=1)[1].expand(-1, -1, h.size(-1)) |
| return torch.gather(h * scores, 1, idx) |
|
|
|
|
| class AASISTPool(nn.Module): |
| def __init__(self, feat_dim, num_layers, gat_dims=(128, 64), |
| pool_ratios=(0.5, 0.7, 0.5, 0.5), temps=(2.0, 2.0, 100.0, 100.0)): |
| super().__init__() |
| g = list(gat_dims) |
| self.proj = nn.Linear(feat_dim, g[0]) |
| self.pos_S = nn.Parameter(torch.randn(1, num_layers, g[0])) |
| self.master1 = nn.Parameter(torch.randn(1, 1, g[0])) |
| self.master2 = nn.Parameter(torch.randn(1, 1, g[0])) |
| self.GAT_S = _GraphAttentionLayer(g[0], g[0], temps[0]) |
| self.GAT_T = _GraphAttentionLayer(g[0], g[0], temps[1]) |
| self.HtrgGAT_11 = _HtrgGraphAttentionLayer(g[0], g[1], temps[2]) |
| self.HtrgGAT_12 = _HtrgGraphAttentionLayer(g[1], g[1], temps[3]) |
| self.HtrgGAT_21 = _HtrgGraphAttentionLayer(g[0], g[1], temps[2]) |
| self.HtrgGAT_22 = _HtrgGraphAttentionLayer(g[1], g[1], temps[3]) |
| self.pool_S = _GraphPool(pool_ratios[0], g[0]) |
| self.pool_T = _GraphPool(pool_ratios[1], g[0]) |
| self.pool_hS1 = _GraphPool(pool_ratios[2], g[1]) |
| self.pool_hT1 = _GraphPool(pool_ratios[2], g[1]) |
| self.pool_hS2 = _GraphPool(pool_ratios[2], g[1]) |
| self.pool_hT2 = _GraphPool(pool_ratios[2], g[1]) |
| self.output_dim = 5 * g[1] |
|
|
| def forward(self, hidden_states): |
| projected = self.proj(torch.stack(hidden_states, dim=1)) |
| e_S = torch.max(torch.abs(projected), dim=2)[0] + self.pos_S |
| e_T = torch.max(torch.abs(projected), dim=1)[0] |
| out_S = self.pool_S(self.GAT_S(e_S)) |
| out_T = self.pool_T(self.GAT_T(e_T)) |
| out_T1, out_S1, m1 = self.HtrgGAT_11(out_T, out_S, self.master1) |
| out_S1, out_T1 = self.pool_hS1(out_S1), self.pool_hT1(out_T1) |
| a_T, a_S, a_m = self.HtrgGAT_12(out_T1, out_S1, m1) |
| out_T1, out_S1, m1 = out_T1 + a_T, out_S1 + a_S, m1 + a_m |
| out_T2, out_S2, m2 = self.HtrgGAT_21(out_T, out_S, self.master2) |
| out_S2, out_T2 = self.pool_hS2(out_S2), self.pool_hT2(out_T2) |
| a_T, a_S, a_m = self.HtrgGAT_22(out_T2, out_S2, m2) |
| out_T2, out_S2, m2 = out_T2 + a_T, out_S2 + a_S, m2 + a_m |
| out_T, out_S, master = torch.max(out_T1, out_T2), torch.max(out_S1, out_S2), torch.max(m1, m2) |
| return torch.cat([torch.max(torch.abs(out_T), 1)[0], out_T.mean(1), |
| torch.max(torch.abs(out_S), 1)[0], out_S.mean(1), master.squeeze(1)], dim=1) |
|
|
|
|
| class MLPBlock(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.net = nn.Sequential(nn.Linear(dim, dim * 2), nn.GELU(), nn.Dropout(0.1), |
| nn.Linear(dim * 2, dim), nn.GELU(), nn.Dropout(0.1)) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class Classifier(nn.Module): |
| def __init__(self, input_dim, hidden_dim, num_layers, num_labels=1): |
| super().__init__() |
| self.input_proj = nn.Linear(input_dim, hidden_dim) if input_dim != hidden_dim else nn.Identity() |
| self.mlp_blocks = nn.ModuleList([MLPBlock(hidden_dim) for _ in range(num_layers)]) |
| self.layer_norms = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(num_layers)]) |
| self.output_head = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.GELU(), |
| nn.Dropout(0.1), nn.Linear(hidden_dim, num_labels)) |
|
|
| def forward(self, x): |
| x = self.input_proj(x) |
| for mlp, ln in zip(self.mlp_blocks, self.layer_norms): |
| x = x + mlp(ln(x)) |
| return self.output_head(x) |
|
|
|
|
| class DeepfakeDetector(nn.Module): |
| def __init__(self, hidden_dim=320, classifier_layers=3, aasist_gat_dims=(128, 64)): |
| super().__init__() |
| self.wavlm = WavLMModel.from_pretrained("microsoft/wavlm-large") |
| self.pool = AASISTPool(self.wavlm.config.hidden_size, self.wavlm.config.num_hidden_layers, aasist_gat_dims) |
| self.classifier = Classifier(self.pool.output_dim, hidden_dim, classifier_layers) |
|
|
| @torch.no_grad() |
| def forward(self, waveform): |
| hs = list(self.wavlm(waveform, output_hidden_states=True).hidden_states[1:]) |
| return self.classifier(self.pool(hs)).squeeze(-1) |
|
|
