File size: 2,999 Bytes
cdfb7d6 5e96656 cdfb7d6 5e96656 cdfb7d6 875a0a4 cdfb7d6 875a0a4 cdfb7d6 5e96656 | 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 | """
Shared NativityClassifier for training and inference.
Variants: MLP (baseline), MLP + channel attention (Squeeze-and-Excitation style).
"""
import torch
import torch.nn as nn
class ImprovedAttentionBlock(nn.Module):
"""
Squeeze-and-Excitation style attention: squeeze 1024 dims to a bottleneck,
then excite back to 1024 channel weights. Learns which channels matter for nativity.
"""
def __init__(self, input_dim=1024, reduction_ratio=8):
super().__init__()
reduced = max(input_dim // reduction_ratio, 32)
self.fc1 = nn.Linear(input_dim, reduced)
self.relu = nn.ReLU()
self.fc2 = nn.Linear(reduced, input_dim)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
# x shape: (batch_size, 1024)
attention_weights = self.sigmoid(self.fc2(self.relu(self.fc1(x))))
return x * attention_weights
class NativityClassifier(nn.Module):
"""
MLP classifier on wav2vec embeddings -> Native vs Non-Native.
Uses BatchNorm and dropout for better generalization.
"""
def __init__(
self,
input_dim=1024,
hidden_dims=(384, 256, 128),
output_dim=2,
dropout=0.35,
):
super(NativityClassifier, self).__init__()
self.input_dim = input_dim
self.output_dim = output_dim
layers = []
prev = input_dim
for h in hidden_dims:
layers.extend([
nn.Linear(prev, h),
nn.BatchNorm1d(h),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
])
prev = h
layers.append(nn.Linear(prev, output_dim))
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)
class NativityClassifierWithAttention(nn.Module):
"""
Squeeze-and-Excitation attention over 1024-d embedding then MLP.
SE block learns which channels matter for non-native prosody; higher final_dropout
reduces overfitting to majority class.
"""
def __init__(
self,
input_dim=1024,
hidden_dims=(384, 256, 128),
output_dim=2,
dropout=0.35,
final_dropout=0.6,
reduction_ratio=8,
):
super(NativityClassifierWithAttention, self).__init__()
self.input_dim = input_dim
self.output_dim = output_dim
self.attention = ImprovedAttentionBlock(input_dim=input_dim, reduction_ratio=reduction_ratio)
layers = []
prev = input_dim
for i, h in enumerate(hidden_dims):
layers.extend([
nn.Linear(prev, h),
nn.BatchNorm1d(h),
nn.ReLU(inplace=True),
nn.Dropout(final_dropout if i == len(hidden_dims) - 1 else dropout),
])
prev = h
layers.append(nn.Linear(prev, output_dim))
self.net = nn.Sequential(*layers)
def forward(self, x):
x = self.attention(x)
return self.net(x)
|