Views
No views yet
.pth (PyTorch).1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4
5class ASR_LSTM_Model(nn.Module):
6 def __init__(self, n_mels=128, vocab_size=36, num_accents=6):
7 super().__init__()
8 self.lstm = nn.LSTM(n_mels, 512, num_layers=4, dropout=0.3, bidirectional=True, batch_first=True)
9 self.ctc_head = nn.Linear(1024, vocab_size)
10 self.accent_head = nn.Linear(1024, num_accents)
11
12 def forward(self, x):
13 out, _ = self.lstm(x)
14 ctc_log_probs = F.log_softmax(self.ctc_head(out), dim=2)
15 pooled_out = out.mean(dim=1)
16 accent_logits = self.accent_head(pooled_out)
17 return ctc_log_probs, accent_logits