altering translate to be compatible with nn.DataParallel

Author: bmccann

OpenNMT/onmt/Models.py

@@ -110,18 +110,18 @@ def forward(self, input, hidden, context, init_output):
  # self.input_feed=False
  outputs = []
  output = init_output
- for emb_t in emb.chunk(emb.size(0)):
+ for i, emb_t in enumerate(emb.chunk(emb.size(0), dim=0)):
  emb_t = emb_t.squeeze(0)
  if self.input_feed:
  emb_t = torch.cat([emb_t, output], 1)
 
- output, hidden = self.rnn(emb_t, hidden)
+ output, h = self.rnn(emb_t, hidden)
  output, attn = self.attn(output, context.t())
  output = self.dropout(output)
  outputs += [output]
 
  outputs = torch.stack(outputs)
- return outputs.transpose(0, 1), hidden, attn
+ return outputs.transpose(0, 1), h, attn
 
 
 class NMTModel(nn.Module):

OpenNMT/onmt/Translator.py

@@ -48,29 +48,32 @@ def buildTargetTokens(self, pred, src, attn):
 
  def translateBatch(self, batch):
  srcBatch, tgtBatch = batch
- batchSize = srcBatch.size(1)
+ batchSize = srcBatch.size(0)
  beamSize = self.opt.beam_size
 
  # (1) run the encoder on the src
 
  # have to execute the encoder manually to deal with padding
  encStates = None
  context = []
- for srcBatch_t in srcBatch.chunk(srcBatch.size(0)):
+ for srcBatch_t in srcBatch.chunk(srcBatch.size(1), dim=1):
  encStates, context_t = self.model.encoder(srcBatch_t, hidden=encStates)
- batchPadIdx = srcBatch_t.data.squeeze(0).eq(onmt.Constants.PAD).nonzero()
+ batchPadIdx = srcBatch_t.data.squeeze(1).eq(onmt.Constants.PAD).nonzero()
  if batchPadIdx.nelement() > 0:
  batchPadIdx = batchPadIdx.squeeze(1)
  encStates[0].data.index_fill_(1, batchPadIdx, 0)
  encStates[1].data.index_fill_(1, batchPadIdx, 0)
  context += [context_t]
 
+ encStates = (self.model._fix_enc_hidden(encStates[0]),
+ self.model._fix_enc_hidden(encStates[1]))
+
  context = torch.cat(context)
  rnnSize = context.size(2)
 
  # This mask is applied to the attention model inside the decoder
  # so that the attention ignores source padding
- padMask = srcBatch.data.eq(onmt.Constants.PAD).t()
+ padMask = srcBatch.data.eq(onmt.Constants.PAD)
  def applyContextMask(m):
  if isinstance(m, onmt.modules.GlobalAttention):
  m.applyMask(padMask)
@@ -85,8 +88,8 @@ def applyContextMask(m):
  initOutput = self.model.make_init_decoder_output(context)
 
  decOut, decStates, attn = self.model.decoder(
- tgtBatch[:-1], decStates, context, initOutput)
- for dec_t, tgt_t in zip(decOut, tgtBatch[1:].data):
+  tgtBatch[:, :-1], decStates, context, initOutput)
+ for dec_t, tgt_t in zip(decOut.transpose(0, 1), tgtBatch.transpose(0, 1)[1:].data):
  gen_t = self.model.generator.forward(dec_t)
  tgt_t = tgt_t.unsqueeze(1)
  scores = gen_t.data.gather(1, tgt_t)
@@ -104,7 +107,7 @@ def applyContextMask(m):
 
  decOut = self.model.make_init_decoder_output(context)
 
- padMask = srcBatch.data.eq(onmt.Constants.PAD).t().unsqueeze(0).repeat(beamSize, 1, 1)
+ padMask = srcBatch.data.eq(onmt.Constants.PAD).unsqueeze(0).repeat(beamSize, 1, 1)
 
  batchIdx = list(range(batchSize))
  remainingSents = batchSize
@@ -117,9 +120,9 @@ def applyContextMask(m):
  if not b.done]).t().contiguous().view(1, -1)
 
  decOut, decStates, attn = self.model.decoder(
- Variable(input), decStates, context, decOut)
+ Variable(input).transpose(0, 1), decStates, context, decOut)
  # decOut: 1 x (beam*batch) x numWords
- decOut = decOut.squeeze(0)
+ decOut = decOut.transpose(0, 1).squeeze(0)
  out = self.model.generator.forward(decOut)
 
  # batch x beam x numWords
@@ -174,7 +177,7 @@ def updateActive(t):
  scores, ks = beam[b].sortBest()
 
  allScores += [scores[:n_best]]
- valid_attn = srcBatch.data[:, b].ne(onmt.Constants.PAD).nonzero().squeeze(1)
+ valid_attn = srcBatch.transpose(0, 1).data[:, b].ne(onmt.Constants.PAD).nonzero().squeeze(1)
  hyps, attn = zip(*[beam[b].getHyp(k) for k in ks[:n_best]])
  attn = [a.index_select(1, valid_attn) for a in attn]
  allHyp += [hyps]
@@ -186,13 +189,14 @@ def translate(self, srcBatch, goldBatch):
  # (1) convert words to indexes
  dataset = self.buildData(srcBatch, goldBatch)
  batch = dataset[0]
+ batch = [x.transpose(0, 1) for x in batch]
 
  # (2) translate
  pred, predScore, attn, goldScore = self.translateBatch(batch)
 
  # (3) convert indexes to words
  predBatch = []
- for b in range(batch[0].size(1)):
+ for b in range(batch[0].size(0)):
  predBatch.append(
  [self.buildTargetTokens(pred[b][n], srcBatch[b], attn[b][n])
  for n in range(self.opt.n_best)]

OpenNMT/train.py

@@ -25,9 +25,9 @@
 
 parser.add_argument('-layers', type=int, default=2,
  help='Number of layers in the LSTM encoder/decoder')
-parser.add_argument('-rnn_size', type=int, default=500,
+parser.add_argument('-rnn_size', type=int, default=512,
  help='Size of LSTM hidden states')
-parser.add_argument('-word_vec_size', type=int, default=500,
+parser.add_argument('-word_vec_size', type=int, default=300,
  help='Word embedding sizes')
 parser.add_argument('-input_feed', type=int, default=1,
  help="""Feed the context vector at each time step as
@@ -43,13 +43,13 @@
 
 ## Optimization options
 
-parser.add_argument('-batch_size', type=int, default=64,
+parser.add_argument('-batch_size', type=int, default=256,
  help='Maximum batch size')
 parser.add_argument('-max_generator_batches', type=int, default=32,
  help="""Maximum batches of words in a sequence to run
  the generator on in parallel. Higher is faster, but uses
  more memory.""")
-parser.add_argument('-epochs', type=int, default=13,
+parser.add_argument('-epochs', type=int, default=50,
  help='Number of training epochs')
 parser.add_argument('-start_epoch', type=int, default=1,
  help='The epoch from which to start')
@@ -58,16 +58,16 @@
  with support (-param_init, param_init)""")
 parser.add_argument('-optim', default='sgd',
  help="Optimization method. [sgd|adagrad|adadelta|adam]")
-parser.add_argument('-learning_rate', type=float, default=1,
+parser.add_argument('-learning_rate', type=float, default=1.0,
  help="""Starting learning rate. If adagrad/adadelta/adam is
  used, then this is the global learning rate. Recommended
  settings: sgd = 1, adagrad = 0.1, adadelta = 1, adam = 0.1""")
 parser.add_argument('-max_grad_norm', type=float, default=5,
  help="""If the norm of the gradient vector exceeds this,
  renormalize it to have the norm equal to max_grad_norm""")
-parser.add_argument('-dropout', type=float, default=0.3,
+parser.add_argument('-dropout', type=float, default=0.2,
  help='Dropout probability; applied between LSTM stacks.')
-parser.add_argument('-learning_rate_decay', type=float, default=0.5,
+parser.add_argument('-learning_rate_decay', type=float, default=0.9,
  help="""Decay learning rate by this much if (i) perplexity
  does not decrease on the validation set or (ii) epoch has
  gone past the start_decay_at_limit""")

OpenNMT/translate.py

@@ -1,7 +1,6 @@
 import onmt
 import torch
 import argparse
-import time
 import math
 
 parser = argparse.ArgumentParser(description='translate.py')
@@ -37,9 +36,11 @@
  help="""If verbose is set, will output the n_best
  decoded sentences""")
 
-parser.add_argument('-cuda', action="store_true",
+parser.add_argument('-gpu', type=int, default=7,
  help="Use CUDA")
 
+
+
 def reportScore(name, scoreTotal, wordsTotal):
  print("%s AVG SCORE: %.4f, %s PPL: %.4f" % (
  name, scoreTotal / wordsTotal,
@@ -48,6 +49,8 @@ def reportScore(name, scoreTotal, wordsTotal):
 
 def main():
  opt = parser.parse_args()
+ opt.cuda = True
+ torch.cuda.set_device(opt.gpu)
 
  translator = onmt.Translator(opt)
 
@@ -58,6 +61,7 @@ def main():
  srcBatch, tgtBatch = [], []
 
  count = 0
+
  tgtF = open(opt.tgt) if opt.tgt else None
  for line in open(opt.src):