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')