File size: 9,745 Bytes
1b92e8f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
import torch
import torch.nn as nn
import torch.nn.functional as F
from atten import Atten


class FGA(nn.Module):
    def __init__(self, vocab_size, word_embed_dim, hidden_ques_dim, hidden_ans_dim,
                 hidden_hist_dim, hidden_cap_dim, hidden_img_dim):
        '''
        Factor Graph Attention
        :param vocab_size: vocabulary size
        :param word_embed_dim
        :param hidden_ques_dim:
        :param hidden_ans_dim:
        :param hidden_hist_dim:
        :param img_features_dim:
        '''
        super(FGA, self).__init__()

        print("Init FGA with vocab size %s, word embed %s, hidden ques %s, hidden ans %s,"
              " hidden hist %s, hidden cap %s, hidden img %s" % (vocab_size, word_embed_dim,
                                                                hidden_ques_dim,
                                                                hidden_ans_dim,
                                                                hidden_hist_dim,
                                                                hidden_cap_dim,
                                                                hidden_img_dim))
        self.hidden_ques_dim = hidden_ques_dim
        self.hidden_ans_dim = hidden_ans_dim
        self.hidden_cap_dim = hidden_cap_dim
        self.hidden_img_dim = hidden_img_dim
        self.hidden_hist_dim = hidden_hist_dim

        # Vocab of History LSTMs is one more as we are keeping a stop id (the last id)
        self.word_embedddings = nn.Embedding(vocab_size+1+1, word_embed_dim, padding_idx=0)

        self.lstm_ques = nn.LSTM(word_embed_dim, self.hidden_ques_dim, batch_first=True)
        self.lstm_ans = nn.LSTM(word_embed_dim, self.hidden_ans_dim, batch_first=True)

        self.lstm_hist_ques = nn.LSTM(word_embed_dim, self.hidden_hist_dim, batch_first=True)
        self.lstm_hist_ans = nn.LSTM(word_embed_dim, self.hidden_hist_dim, batch_first=True)

        self.lstm_hist_cap = nn.LSTM(word_embed_dim, self.hidden_cap_dim, batch_first=True)


        self.qahistnet = nn.Sequential(
            nn.Linear(self.hidden_hist_dim*2, self.hidden_hist_dim),
            nn.ReLU(inplace=True)
        )

        self.concat_dim = self.hidden_ques_dim + self.hidden_ans_dim + \
                          self.hidden_ans_dim + self.hidden_img_dim + \
                          self.hidden_cap_dim + self.hidden_hist_dim*9

        self.simnet = nn.Sequential(
            nn.Linear(self.concat_dim, (self.concat_dim)//2, bias=False),
            nn.BatchNorm1d((self.concat_dim) // 2),
            nn.ReLU(inplace=True),
            nn.Linear((self.concat_dim)//2, (self.concat_dim)//4, bias=False),
            nn.BatchNorm1d((self.concat_dim) // 4),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear((self.concat_dim)//4, 1)
        )

        # To share weights, provide list of tuples: (idx, list of connected utils)
        # Note, for efficiency, the shared utils (i.e., history, are connected to ans and question only.
        # connecting shared factors is not supported (!)
        sharing_factor_weights = {4: (9, [0, 1]),
                                  5: (9, [0, 1])}

        self.mul_atten = Atten(util_e=[self.hidden_ans_dim, # Answer modal
                                       self.hidden_ques_dim, # Question modal
                                       self.hidden_cap_dim, # Caption modal
                                       self.hidden_img_dim, # Image modal
                                       self.hidden_hist_dim, # Question-history modal
                                       self.hidden_hist_dim # Answer-history modal
                                       ],
                               sharing_factor_weights=sharing_factor_weights,
                               sizes=[100, # 100 Answers
                                      21, # Question length
                                      41, # Caption length
                                      37, # 36 Image regions
                                      21, # History-Question length
                                      21 #  History-Answer length
                                      ] # The spatial dim used for pairwise normalization (use force for adaptive)
                               , prior_flag=True,
                               pairwise_flag=True)



    def forward(self, input_ques, input_ans, input_hist_ques, input_hist_ans, input_hist_cap,
                input_ques_length, input_ans_length, input_cap_length, i_e):
        """

        :param input_ques:
        :param input_ans:
        :param input_hist_ques:
        :param input_hist_ans:
        :param input_hist_cap:
        :param input_ques_length:
        :param input_ans_length:
        :param input_cap_length:
        :param i_e:
        :return:
        """


        n_options = input_ans.size()[1]
        batch_size = input_ques.size()[0]



        nqa_per_dial, nwords_per_qa = input_hist_ques.size()[1], input_hist_ques.size()[2]
        nwords_per_cap = input_hist_cap.size()[1]
        max_length_input_ans = input_ans.size()[-1]

        assert batch_size == input_hist_ques.size()[0] == input_hist_ans.size()[0] == input_ques.size()[0] == \
               input_ans.size()[0] == input_hist_cap.size()[0]
        assert nqa_per_dial == input_hist_ques.size()[1] == input_hist_ans.size()[1]
        assert nwords_per_qa == input_hist_ques.size()[2] == input_hist_ans.size()[2]

        q_we = self.word_embedddings(input_ques)
        a_we = self.word_embedddings(input_ans.view(-1, max_length_input_ans))
        hq_we = self.word_embedddings(input_hist_ques.view(-1, nwords_per_qa))
        ha_we = self.word_embedddings(input_hist_ans.view(-1, nwords_per_qa))
        c_we = self.word_embedddings(input_hist_cap.view(-1, nwords_per_cap))



        '''
            q_we = batch x 20 x embed_ques_dim
            a_we = 100*batch x 20 x embed_ans_dim
            hq_we = batch*nqa_per_dial, nwords_per_qa, embed_hist_dim
            ha_we = batch*nqa_per_dial, nwords_per_qa, embed_hist_dim
            c_we = batch*ncap_per_dial, nwords_per_cap, embed_hist_dim
        '''
        self.lstm_ques.flatten_parameters()
        self.lstm_ans.flatten_parameters()
        self.lstm_hist_ques.flatten_parameters()
        self.lstm_hist_ans.flatten_parameters()
        self.lstm_hist_cap.flatten_parameters()


        i_feat = i_e

        q_seq, self.hidden_ques = self.lstm_ques(q_we)
        a_seq, self.hidden_ans = self.lstm_ans(a_we)
        hq_seq, self.hidden_hist_ques = self.lstm_hist_ques(hq_we)
        ha_seq, self.hidden_hist_ans = self.lstm_hist_ans(ha_we)
        cap_seq, self.hidden_cap = self.lstm_hist_cap(c_we)


        '''
            length is used for attention prior
        '''
        q_len = input_ques_length.data - 1
        c_len = input_cap_length.data.view(-1) - 1


        ans_index = torch.arange(0, n_options * batch_size).long().cuda()
        ans_len = input_ans_length.data.view(-1) - 1
        ans_seq = a_seq[ans_index, ans_len, :]
        ans_seq = ans_seq.view(batch_size, n_options, self.hidden_ans_dim)

        batch_index = torch.arange(0, batch_size).long().cuda()
        q_prior = torch.zeros(batch_size, q_seq.size(1)).cuda()
        q_prior[batch_index, q_len] = 100
        c_prior = torch.zeros(batch_size, cap_seq.size(1)).cuda()
        c_prior[batch_index, c_len] = 100
        ans_prior = torch.ones(batch_size, ans_seq.size(1)).cuda()
        img_prior = torch.ones(batch_size, i_feat.size(1)).cuda()

        (ans_atten, ques_atten, cap_atten, img_atten, hq_atten, ha_atten) = \
            self.mul_atten([ans_seq, q_seq, cap_seq, i_feat, hq_seq, ha_seq],
                           priors=[ans_prior, q_prior, c_prior, img_prior, None, None])

        '''
            expand to answers based
        '''
        ques_atten = torch.unsqueeze(ques_atten, 1).expand(batch_size,
                                                           n_options,
                                                           self.hidden_ques_dim)
        cap_atten = torch.unsqueeze(cap_atten, 1).expand(batch_size,
                                                         n_options,
                                                         self.hidden_cap_dim)
        img_atten = torch.unsqueeze(img_atten, 1).expand(batch_size, n_options,
                                                         self.hidden_img_dim)
        ans_atten = torch.unsqueeze(ans_atten, 1).expand(batch_size, n_options,
                                                         self.hidden_ans_dim)


        '''
            combine history
        '''

        input_qahistnet = torch.cat((hq_atten, ha_atten), 1)
        # input_qahistnet: (nqa_per_dial*batch x 2*hidden_hist_dim)
        output_qahistnet = self.qahistnet(input_qahistnet)
        # output_qahistnet: (nqa_per_dial*batch x hidden_hist_dim)
        output_qahistnet = output_qahistnet.view(batch_size,
                                                 nqa_per_dial * self.hidden_hist_dim)
        # output_qahistnet: (batch x nqa_per_dial*hidden_hist_dim)
        output_qahistnet = torch.unsqueeze(output_qahistnet, 1)\
            .expand(batch_size,
                    n_options,
                    nqa_per_dial * self.hidden_hist_dim)

        input_qa = torch.cat((ans_seq, ques_atten, ans_atten, img_atten,
                              output_qahistnet, cap_atten), 2)  # Concatenate last dimension

        input_qa = input_qa.view(batch_size * n_options, self.concat_dim)

        out_scores = self.simnet(input_qa)

        out_scores = out_scores.squeeze(dim=1)
        out_scores = out_scores.view(batch_size, n_options)

        return out_scores