File size: 642 Bytes
bc45d7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | from dataclasses import asdict
import torch
import torch.nn as nn
import torch.nn.functional as F
class LSTM(nn.Module):
def __init__(self, config, n_classes=50):
super().__init__()
config_dict = asdict(config)
self.lstm = nn.LSTM(**config_dict)
in_features = (
config.hidden_size * 2 if config.bidirectional else config.hidden_size
)
self.l1 = nn.Linear(in_features=in_features, out_features=n_classes)
def forward(self, x):
x, (_, _) = self.lstm(x)
x = torch.max(x, dim=1).values
x = F.dropout(x, p=0.3)
x = self.l1(x)
return x
|