|
- import numpy
-
- import torch
- import torch.nn as nn
- from torch.autograd import Variable
- import torch.nn.functional as F
-
- class HAN(nn.Module):
- def __init__(self, input_size, output_size,
- word_hidden_size, word_num_layers, word_context_size,
- sent_hidden_size, sent_num_layers, sent_context_size):
- super(HAN, self).__init__()
-
- self.word_layer = AttentionNet(input_size,
- word_hidden_size,
- word_num_layers,
- word_context_size)
- self.sent_layer = AttentionNet(2* word_hidden_size,
- sent_hidden_size,
- sent_num_layers,
- sent_context_size)
- self.output_layer = nn.Linear(2* sent_hidden_size, output_size)
- self.softmax = nn.LogSoftmax(dim=1)
-
- def forward(self, batch_doc):
- # input is a sequence of vector
- # if level == w, a seq of words (a sent); level == s, a seq of sents (a doc)
- doc_vec_list = []
- for doc in batch_doc:
- s_list = []
- for sent in doc:
- s_list.append(self.word_layer(sent))
- s_vec = torch.cat(s_list, dim=0)
- vec = self.sent_layer(s_vec)
- doc_vec_list.append(vec)
- doc_vec = torch.cat(doc_vec_list, dim=0)
- output = self.softmax(self.output_layer(doc_vec))
- return output
-
- class AttentionNet(nn.Module):
- def __init__(self, input_size, gru_hidden_size, gru_num_layers, context_vec_size):
- super(AttentionNet, self).__init__()
-
- self.input_size = input_size
- self.gru_hidden_size = gru_hidden_size
- self.gru_num_layers = gru_num_layers
- self.context_vec_size = context_vec_size
-
- # Encoder
- self.gru = nn.GRU(input_size=input_size,
- hidden_size=gru_hidden_size,
- num_layers=gru_num_layers,
- batch_first=False,
- bidirectional=True)
- # Attention
- self.fc = nn.Linear(2* gru_hidden_size, context_vec_size)
- self.tanh = nn.Tanh()
- self.softmax = nn.Softmax(dim=0)
- # context vector
- self.context_vec = nn.Parameter(torch.Tensor(context_vec_size, 1))
- self.context_vec.data.uniform_(-0.1, 0.1)
-
- def forward(self, inputs):
- # inputs's dim (seq_len, word_dim)
- inputs = torch.unsqueeze(inputs, 1)
- h_t, hidden = self.gru(inputs)
- h_t = torch.squeeze(h_t, 1)
- u = self.tanh(self.fc(h_t))
- alpha = self.softmax(torch.mm(u, self.context_vec))
- output = torch.mm(h_t.t(), alpha).t()
- # output's dim (1, 2*hidden_size)
- return output
-
|