import torch import torch.nn as nn import math import librosa class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads self.W_q = nn.Linear(d_model, d_model) # query self.W_k = nn.Linear(d_model, d_model) # key self.W_v = nn.Linear(d_model, d_model) # value self.W_o = nn.Linear(d_model, d_model) # output def scaled_dot_product_attention(self, Q, K, V, mask=None): attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: attn_scores = attn_scores.masked_fill(mask == 0, -1e9) attn_probs = torch.softmax(attn_scores, dim=-1) output = torch.matmul(attn_probs, V) return output def split_heads(self, x): batch_size, seq_length, d_model = x.size() return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2) def combine_heads(self, x): batch_size, _, seq_length, d_k = x.size() return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model) def forward(self, Q, K, V, mask=None): Q = self.split_heads(self.W_q(Q)) K = self.split_heads(self.W_k(K)) V = self.split_heads(self.W_v(V)) attn_output = self.scaled_dot_product_attention(Q, K, V, mask) output = self.W_o(self.combine_heads(attn_output)) return output class PositionWiseFeedForward(nn.Module): def __init__(self, d_model, d_ff): super(PositionWiseFeedForward, self).__init__() self.fc1 = nn.Linear(d_model, d_ff) self.fc2 = nn.Linear(d_ff, d_model) self.relu = nn.ReLU() def forward(self, x): return self.fc2(self.relu(self.fc1(x))) class PositionalEncoding(nn.Module): def __init__(self, d_model, max_seq_length): super(PositionalEncoding, self).__init__() pe = torch.zeros(max_seq_length, d_model) position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer('pe', pe.unsqueeze(0)) def forward(self, x): self.pe = self.pe.to(x.device) return x + self.pe[:, :x.size(1)] class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super(EncoderLayer, self).__init__() self.self_attn = MultiHeadAttention(d_model, num_heads) self.feed_forward = PositionWiseFeedForward(d_model, d_ff) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask): attn_output = self.self_attn(x, x, x, mask) x = self.norm1(x + self.dropout(attn_output)) ff_output = self.feed_forward(x) x = self.norm2(x + self.dropout(ff_output)) return x class PhantomNet(nn.Module): def __init__(self, use_mode, feature_size, conv_projection, num_classes, num_heads=8, num_layers=6, d_ff=2048, dropout=0.1): super(PhantomNet, self).__init__() self.conv1 = nn.Conv1d(in_channels=1, out_channels=512, kernel_size=10, stride=5) self.conv2 = nn.Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2) self.conv3 = nn.Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2) self.conv4 = nn.Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2) self.conv5 = nn.Conv1d(in_channels=512, out_channels=512, kernel_size=3, stride=2) self.conv6 = nn.Conv1d(in_channels=512, out_channels=512, kernel_size=2, stride=2) self.conv7 = nn.Conv1d(in_channels=512, out_channels=512, kernel_size=2, stride=2) self.use_mode = use_mode self.conv_projection = conv_projection self.num_classes = num_classes self.flatten = nn.Flatten() self.sigmoid = nn.Sigmoid() self.gelu = nn.GELU() self.relu = nn.ReLU() self.fcIntermidiate = nn.Linear(512, feature_size) self.positional_encoding = PositionalEncoding(feature_size, 10000) self.encoder_layers = nn.ModuleList( [EncoderLayer(feature_size, num_heads, d_ff, dropout) for _ in range(num_layers)]) self.dropout = nn.Dropout(dropout) if self.conv_projection: self.convProjection = nn.Conv1d(feature_size, feature_size, kernel_size=128, stride=1) self.fc1 = nn.Linear(feature_size, feature_size) self.fc2 = nn.Linear(feature_size, 1, bias=True) if self.use_mode == 'spoof': #if there is a mismatch error, you will need to replace this input size.. currently working with 8 seconds samples #just multiply 95.760 * seconds the get this layer's input size #or I can just add another parameter to the model seq_length and input = seq_length * feature_size self.fcSpoof = nn.Linear(286080, d_ff) self.fcFinal = nn.Linear(d_ff,self.num_classes) else: self.fcSpoof = None def forward(self, src): src = src.unsqueeze(1) src = self.gelu(self.conv1(src)) src = self.gelu(self.conv2(src)) src = self.gelu(self.conv3(src)) src = self.gelu(self.conv4(src)) src = self.gelu(self.conv5(src)) src = self.gelu(self.conv6(src)) src = self.gelu(self.conv7(src)) src = src.permute(0, 2, 1) src = self.fcIntermidiate(src) src = src.permute(0, 2, 1) if self.conv_projection: src = self.gelu(self.convProjection(src)) src = self.dropout(src) src = src.transpose(1, 2) src_embedded = self.dropout(self.positional_encoding(src)) enc_output = src_embedded for enc_layer in self.encoder_layers: enc_output = enc_layer(enc_output, None) embeddings = self.fc1(enc_output) flatten_embeddings = self.flatten(embeddings) if self.use_mode == 'extractor': return embeddings elif self.use_mode == 'partialSpoof': return self.fc2(embeddings) elif self.use_mode == 'spoof': out_fcSpoof= self.fcSpoof(flatten_embeddings) output = self.fcFinal(out_fcSpoof) # output = self.sigmoid(self.fcSpoof(flatten_embeddings)) # print(f"Model output shape: {output.shape}") return output else: raise ValueError('Wrong use mode of PhantomNet, please pick between extractor, partialSpoof, or spoof')