reverting cudnn decoder to lstmcell

OpenNMT/onmt/Models.py

@@ -38,6 +38,34 @@ def forward(self, input, hidden=None):
         return hidden_t, outputs
 
 
+class StackedLSTM(nn.Module):
+    def __init__(self, num_layers, input_size, rnn_size, dropout):
+        super(StackedLSTM, self).__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.num_layers = num_layers
+        self.layers = nn.ModuleList()
+
+        for i in range(num_layers):
+            self.layers.append(nn.LSTMCell(input_size, rnn_size))
+            input_size = rnn_size
+
+    def forward(self, input, hidden):
+        h_0, c_0 = hidden
+        h_1, c_1 = [], []
+        for i, layer in enumerate(self.layers):
+            h_1_i, c_1_i = layer(input, (h_0[i], c_0[i]))
+            input = h_1_i
+            if i != self.num_layers:
+                input = self.dropout(input)
+            h_1 += [h_1_i]
+            c_1 += [c_1_i]
+
+        h_1 = torch.stack(h_1)
+        c_1 = torch.stack(c_1)
+
+        return input, (h_1, c_1)
+
+
 class Decoder(nn.Module):
 
     def __init__(self, opt, dicts):
@@ -51,9 +79,7 @@ def __init__(self, opt, dicts):
         self.word_lut = nn.Embedding(dicts.size(),
                                   opt.word_vec_size,
                                   padding_idx=onmt.Constants.PAD)
-        self.rnn = nn.LSTM(input_size, opt.rnn_size, 
-                        num_layers=opt.layers,
-                        dropout=opt.dropout)
+        self.rnn = StackedLSTM(opt.layers, input_size, opt.rnn_size, opt.dropout)
         self.attn = onmt.modules.GlobalAttention(opt.rnn_size)
         self.dropout = nn.Dropout(opt.dropout)
 
@@ -77,16 +103,16 @@ def forward(self, input, hidden, context, init_output):
         outputs = []
         output = init_output
         for i, emb_t in enumerate(emb.split(1)):
-            emb_t = emb_t
+            emb_t = emb_t.squeeze(0)
             if self.input_feed:
-                emb_t = torch.cat([emb_t, output], 2)
+                emb_t = torch.cat([emb_t, output], 1)
 
             output, hidden = self.rnn(emb_t, hidden)
-            output, attn = self.attn(output.squeeze(0), context.t())
-            output = self.dropout(output.unsqueeze(0))
+            output, attn = self.attn(output, context.t())
+            output = self.dropout(output)
             outputs += [output]
 
-        outputs = torch.cat(outputs, 0)
+        outputs = torch.stack(outputs)
         return outputs.transpose(0, 1), hidden, attn
 
 
@@ -105,7 +131,7 @@ def set_generate(self, enabled):
     def make_init_decoder_output(self, context):
         batch_size = context.size(1)
         h_size = (batch_size, self.decoder.hidden_size)
-        return Variable(context.data.new(1, *h_size).zero_(), requires_grad=False)
+        return Variable(context.data.new(*h_size).zero_(), requires_grad=False)
 
     def _fix_enc_hidden(self, h):
         #  the encoder hidden is  (layers*directions) x batch x dim