Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
253 lines (207 sloc) 10.4 KB
import os
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
class QuestionModule(nn.Module):
def __init__(self, vocab_size, hidden_size):
super(QuestionModule, self).__init__()
self.vocab_size = vocab_size # Size of the vocabulary used in word embedding
self.hidden_size = hidden_size # Size of the hidden state of GRU
self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True)
def forward(self, questions, word_embedding):
# questions.size() = (batch_size, num_tokens)
# word_embedding -> (batch_size, num_tokens, embedding_length)
# self.gru() -> (1, batch_size, hidden_size)
questions = word_embedding(questions) # Word embedding of the question
output, questions = self.gru(questions) # What is the initial hidden vector given to GRU?
questions = torch.transpose(questions, 0, 1)
return questions
class InputModule(nn.Module):
def __init__(self, vocab_size, hidden_size):
super(InputModule, self).__init__()
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.gru = nn.GRU(hidden_size, hidden_size, bidirectional=True, batch_first=True)
for name, param in self.gru.state_dict().items():
if 'weight' in name:
init.xavier_normal(param)
self.dropout = nn.Dropout(0.1)
''' We will now define the encoding scheme which is positional encoding in the paper " Dynamic Memory Network for Textual and Visual
Question Answering '''
def positional_encoder(embedded_sentence):
# embedded_sentence.size() = (batch_size, num_sentences, num_tokens, embedding_length)
# l.size() = (num_tokems, embedding_length)
# output.size() = (num_batch, num_sentences, embedding_length)
# The outputs are basically f1, f2, f3,.... which will go into the input fusion layer in the next step to add share information
# between sentences using a BiDirfectional GRU module.
batch_size, num_sentences, num_tokens, embedding_length = embedded_sentence.size()
l = [] # It will be same for all sentences in all batches as num_tokens and embedding_length is same for the entire dataset.
for j in range(num_tokens):
x = []
for d in range(embedding_length):
x.append((1 - (j/(num_tokens-1))) - (d/(embedding_length-1)) * (1 - 2*j/(num_tokens-1)))
l.append(x)
l = torch.FloatTensor(l)
l = l.unsqueeze(0) # adding an extra dimension at first place for batch_size
l = l.unsqueeze(1) # adding an extra dimension at sencond place for num_sentences
l = l.expand_as(embedded_sentence) # so that l.size() = (batch_size, num_sentences, num_tokens, embedding_length)
mat = embedded_sentence*Variable(l.cuda())
f_ids = torch.sum(mat, dim=2).squeeze(2) # sum along token dimension
return f_ids
def forward(self, input, word_embedding):
# input.size() = (batch_size, num_sentences, num_tokens)
# word_embedding -> (batch_size, num_sentences, num_tokens, embedding_length)
# positional_encoder(word_embedding(input)) -> (batch_size, num_sentences, embedding_length)
# Now BidirectionalGRU blocks receive their input, the output of the positional encoder and finally give facts
# facts.size() = (batch_size, num_sentences, embedding_length) embedding_length = hidden_size
input = input.view(input.size()[0], -1)# Isn't it already in this format ?
input = word_embedding(input)
input = input.view(input.size()[0], input.size()[1], input.size()[2], -1)
input = self.positional_encoder(input)
input = self.dropout(input)
h0 = Variable(torch.zeros(2, input.size()[0], self.hidden_size).cuda()) # Initializing the initial hidden state (at t=0 time step)
facts, hdn = self.gru(input, h0)
facts = facts[:, :, :hidden_size] + facts[:, :, hidden_size:]
return facts
class AttnGRUCell(nn.Module):
def __init__(self, input_size, hidden_size):
super(AttnGRUCell, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.Wr = nn.Linear(input_size, hidden_size)
self.Ur = nn.Linear(hidden_size, hidden_size)
self.W = nn.Linear(input_size, hidden_size)
self.U = nn.Linear(hidden_size, hidden_size)
init.xavier_normal(self.Wr.state_dict()['weight'])
init.xavier_normal(self.Ur.state_dict()['weight'])
init.xavier_normal(self.W.state_dict()['weight'])
init.xavier_normal(self.U.state_dict()['weight'])
def forward(self, fact, hi_1, g):
# fact is the final output of InputModule for each sentence and fact.size() = (batch_size, embedding_length)
# hi_1.size() = (batch_size, embedding_length=hidden_size)
# g.size() = (batch_size, )
r_i = F.sigmoid(self.Wr(fact) + self.Ur(hi_1))
h_tilda = F.tanh(self.W(fact) + r*self.U(hi_1))
hi = g*h_tilda + (1 - g)*hi_1
return hi # Returning the next hidden state considering the first fact and so on.
class AttnGRU(nn.Module):
def __init__(self, input_size, hidden_size):
super(AttnGRU, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.AttnGRUCell = AttnGRUCell(input_size, hidden_size)
def forward(self, facts, G):
# facts.size() = (batch_size, num_sentences, embedding_length)
# fact.size() = (batch_size, embedding_length=hidden_size)
# G.size() = (batch_size, num_sentences)
# g.size() = (batch_size, )
h_0 = Variable(torch.zeros(self.hidden_size)).cuda()
for sen in range(facts.size()[1]):
fact = facts[:, sen, :]
g = G[:, sen]
if sen == 0: # Initialization for first sentence only
hi_1 = h_0.unsqueeze(0).expand_as(fact)
hi_1 = self.AttnGRUCell(fact, hi_1, g)
C = hi_1 # Final hidden vector as the contextual vector used for updating memory
return C
class MemoryModule(nn.Module): # Takes facts, question and prev_mem as its and output next_mem
def __init__(self, hidden_size):
super(MemoryModule, self).__init__()
self.hidden_size = hidden_size
self.AttnGRU = AttnGRU(hidden_size, hidden_size)
self.W1 = nn.Linear(4*hidden_size, hidden_size)
self.W2 = nn.Linear(hidden_size, 1)
self.W_mem = nn.Linear(3*hidden_size, hidden_size)
init.xavier_normal(self.W1.state_dict()['weight'])
init.xavier_normal(self.W2.state_dict()['weight'])
init.xavier_normal(self.W_mem.state_dict()['weight'])
def gateMatrix(self, facts, questions, prev_mem):
# facts.size() = (batch_size, num_sentences, embedding_length=hidden_size)
# questions.size() = (batch_size, 1, embedding_length)
# prev_mem.size() = (batch_size, 1, embedding_length)
# z.size() = (batch_size, num_sentences, 4*embedding_length)
# G.size() = (batch_size, num_sentences)
questions = questions.expand_as(facts)
prev_mem = prev_mem.expand_as(facts)
z = torch.cat([facts*questions, facts*prev_mem, torch.abs(facts - questions), torch.abs(facts - prev_mem)], dim=2)
# z.size() = (batch_size, num_sentences, 4*embedding_length)
z = z.view(-1, 4*embedding_length)
Z = self.W2(F.tanh(self.W1(z)))
Z = Z.view(batch_size, -1)
G = F.softmax(Z)
return G
def forward(self, facts, questions, prev_mem):
# questions = questions.unsqueeze(1)
# prev_mem = prev_mem.unsqueeze(1)
G = self.gateMatrix(facts, questions, prev_mem)
C = self.AttnGRU(facts, G)
# Now considering prev_mem, C and question, we will update the memory state as follows
concat = torch.cat([prev_mem.squeeze(1), C, questions.squeeze(1)], dim=1)
next_mem = F.relu(self.W_mem(concat))
next_mem = next_mem.unsqueeze(1)
return next_mem
class AnswerModule(nn.Module):
def __init__(self, vocab_size, hidden_size):
super(AnswerModule, self).__init__()
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.W = nn.Linear(2*hidden_size, vocab_size)
init.xavier_normal(self.W.state_dict()['weight'])
self.dropout = nn.Dropout(0.1)
def forward(self, final_mem, questions):
final_mem = self.dropout(final_mem)
concat = torch.cat([final_mem, questions], dim=2).squeeze(1)
out = self.W(concat) # As per the paper, we are concatenating the final memory state m_T, and the question q and passing
# this resultant vector to a linear layer
return out
''' We define the model for the network incorporating the input, question, answer and the episodic memory module. We use the Cross Entropy loss criterion for measuring loss'''
class DMN(nn.Module):
def __init__(self, hidden_size, vocab_size, num_pass=3, qa=None):
super(DMN,self).__init__()
self.num_pass= num_pass
self.qa= qa
self.word_embedding= nn.Embedding(vocab_size, hidden_size, padding_idx=0, sparse=True)
init.uniform(self.word_embedding.state_dict()['weight'], a= -(3**0.5), b=3**0.5)
self.criterion= nn.CrossEntropyLoss(size_average=False)
self.input_module= InputModule(vocab_size,hidden_size) ##Vocab size refers to the size of vocabulary used
self.question_module= QuestionModule(vocab_size, hidden_size)
self.memory= MemoryModule(hidden_size)
self.answer_module= AnswerModule(vocab_size,hidden_size)
def forward(self, context, questions):
#facts.size()= (batch_size, num_sentences, embedding_length= hidden.size())
#questions.size() = (batch_size, 1, embedding_length)
facts= self.input_module(context, self.word_embedding)
questions= self.question_module(questions, self.word_embedding)
X= questions
for passes in range(self.num_pass):
X= self.memory(facts, questions, X)
pred= self.answer_module(X, questions)
return pred_id
'''Total loss to be calculated '''
def loss(self,context, questions, targets):
output= self.forward(context, questions)
loss= self.criterion(output, targets)
para_loss= 0
for param in self.parameters():
para_loss+= 0.001* torch.sum(param*param)
pred= F.softmax(output)
_, pred_id= torch.max(pred, dim=1)
correct= (pred_id.data == answers.data)
acc= torch.mean(correct.float())
return loss+para_loss, acc
def interpret_indexed_tensor(self,var):
if len(var.size()) == 3:
for n, sentences in enumerate(var):
s= ' '.join([self.qa.IVOCAB[elem.data[0]] for elem in sentence])
print (str(n)+'th batch, '+str(i)+'th sentence, '+str(s))
elif len(var.size()) == 2:
for n, sentence in enumerate(var):
s= ' '.join([self.qa.IVOCAB[elem.data[0]] for elem in sentence])
print (str(n)+'th batch, '+str(s))
elif len(var.size()) == 1:
for n, token in enumerate(var):
s= self.qa.IVOCAB[token.data[0]]
print (str(n)+'th of batch, '+str(s))
You can’t perform that action at this time.