-
Notifications
You must be signed in to change notification settings - Fork 8
/
lstm.py
24 lines (20 loc) · 845 Bytes
/
lstm.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# -*- coding: utf-8 -*-
# file: lstm.py
# author: songyouwei <youwei0314@gmail.com>
# Copyright (C) 2018. All Rights Reserved.
from layers.dynamic_rnn import DynamicLSTM
import torch
import torch.nn as nn
class LSTM(nn.Module):
def __init__(self, embedding_matrix, opt):
super(LSTM, self).__init__()
self.embed = nn.Embedding.from_pretrained(torch.tensor(embedding_matrix, dtype=torch.float))
self.lstm = DynamicLSTM(opt.embed_dim, opt.hidden_dim, num_layers=1, batch_first=True)
self.dense = nn.Linear(opt.hidden_dim, opt.polarities_dim)
def forward(self, inputs):
text_raw_indices = inputs[0]
x = self.embed(text_raw_indices)
x_len = torch.sum(text_raw_indices != 0, dim=-1)
_, (h_n, _) = self.lstm(x, x_len.cpu())
out = self.dense(h_n[0])
return out