File size: 4,207 Bytes
aee40b4 | 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 | import torch
import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
class BiLSTMFeatureExtractor(nn.Module):
"""
BiLSTM Feature Extractor module.
Inputs:
x : Padded WavLM sequence embeddings [B, T, 768]
mask : Timestep mask [B, T] (optional, 1 for real timesteps, 0 for padding)
Outputs:
out : BiLSTM contextual representations [B, T, 256]
"""
def __init__(self, input_size=768, hidden_size=128, num_layers=1, dropout=0.0):
super(BiLSTMFeatureExtractor, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.num_layers = num_layers
self.bidirectional = True
# BiLSTM Layer: hidden_size=128, bidirectional=True -> Output dim = 128 * 2 = 256
self.bilstm = nn.LSTM(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
bidirectional=True,
dropout=dropout if num_layers > 1 else 0.0
)
def forward(self, x, mask=None):
batch_size, seq_len, _ = x.shape
if mask is not None:
lengths = mask.sum(dim=1).cpu()
packed_x = pack_padded_sequence(
x,
lengths,
batch_first=True,
enforce_sorted=False
)
packed_out, (hn, cn) = self.bilstm(packed_x)
out, _ = pad_packed_sequence(
packed_out,
batch_first=True,
total_length=seq_len
)
else:
out, (hn, cn) = self.bilstm(x)
return out
class TemporalAttention(nn.Module):
"""
Temporal Attention Module.
Learns frame-wise attention scores over sequence length T to aggregate
BiLSTM output sequence [B, T, 256] into a fixed 256-dimensional context vector [B, 256].
Inputs:
h : BiLSTM outputs [B, T, input_dim] (input_dim=256)
mask : Timestep mask [B, T] (1 for real timesteps, 0 for padding)
Outputs:
context : Aggregated context vector [B, 256]
attn_weights : Normalized attention weights [B, T]
"""
def __init__(self, input_dim=256):
super(TemporalAttention, self).__init__()
self.input_dim = input_dim
self.w = nn.Linear(input_dim, 1, bias=False)
def forward(self, h, mask=None):
scores = self.w(torch.tanh(h)).squeeze(-1) # [B, T]
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn_weights = torch.softmax(scores, dim=1) # [B, T]
context = torch.bmm(attn_weights.unsqueeze(1), h).squeeze(1) # [B, 256]
return context, attn_weights
class BiLSTMAttentionClassifier(nn.Module):
"""
Full Downstream Model Architecture:
WavLM Embedding [B, T, 768] -> BiLSTM -> [B, T, 256] -> Temporal Attention -> [B, 256] -> Linear Classifier -> [B, 6] Logits
"""
def __init__(self, input_size=768, hidden_size=128, num_classes=6, dropout=0.3):
super(BiLSTMAttentionClassifier, self).__init__()
self.bilstm = BiLSTMFeatureExtractor(
input_size=input_size,
hidden_size=hidden_size,
num_layers=1,
dropout=dropout
)
context_dim = hidden_size * 2 # 256
self.attention = TemporalAttention(input_dim=context_dim)
self.dropout = nn.Dropout(dropout)
self.classifier = nn.Linear(context_dim, num_classes)
def forward(self, x, mask=None):
# 1. BiLSTM: [B, T, 768] -> [B, T, 256]
bilstm_out = self.bilstm(x, mask=mask)
# 2. Temporal Attention: [B, T, 256] -> [B, 256]
context, attn_weights = self.attention(bilstm_out, mask=mask)
# 3. Dropout + Linear Classifier: [B, 256] -> [B, 6]
dropped_context = self.dropout(context)
logits = self.classifier(dropped_context)
return logits, attn_weights
|