File size: 7,763 Bytes
67982e2 | 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """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): # (B, T) 16 kHz -> logit (B,); sigmoid(logit)=P(real)
hs = list(self.wavlm(waveform, output_hidden_states=True).hidden_states[1:])
return self.classifier(self.pool(hs)).squeeze(-1)
|