new DataParallel allows dim 1; remove unnecessary transposes; add train_ppl to chkpt

Author: bmccann

OpenNMT/onmt/Models.py

@@ -26,9 +26,10 @@ def __init__(self, opt, dicts):
  self.word_lut.weight.copy_(pretrained)
 
  def forward(self, input, hidden=None):
- batch_size = input.size(0) # batch first for multi-gpu compatibility
- emb = self.word_lut(input).transpose(0, 1)
+ emb = self.word_lut(input)
+
  if hidden is None:
+ batch_size = emb.size(1)
  h_size = (self.layers * self.num_directions, batch_size, self.hidden_size)
  h_0 = Variable(emb.data.new(*h_size).zero_(), requires_grad=False)
  c_0 = Variable(emb.data.new(*h_size).zero_(), requires_grad=False)
@@ -91,9 +92,9 @@ def __init__(self, opt, dicts):
 
 
  def forward(self, input, hidden, context, init_output):
- emb = self.word_lut(input).transpose(0, 1)
+ emb = self.word_lut(input)
 
- batch_size = input.size(0)
+ batch_size = input.size(1)
 
  h_size = (batch_size, self.hidden_size)
 
@@ -102,7 +103,7 @@ def forward(self, input, hidden, context, init_output):
  # self.input_feed=False
  outputs = []
  output = init_output
- for i, emb_t in enumerate(emb.split(1)):
+ for emb_t in emb.split(1):
  emb_t = emb_t.squeeze(0)
  if self.input_feed:
  emb_t = torch.cat([emb_t, output], 1)
@@ -113,7 +114,7 @@ def forward(self, input, hidden, context, init_output):
  outputs += [output]
 
  outputs = torch.stack(outputs)
- return outputs.transpose(0, 1), hidden, attn
+ return outputs, hidden, attn
 
 
 class NMTModel(nn.Module):
@@ -145,7 +146,7 @@ def _fix_enc_hidden(self, h):
 
  def forward(self, input):
  src = input[0]
- tgt = input[1][:, :-1] # exclude last target from inputs
+ tgt = input[1][:-1] # exclude last target from inputs
  enc_hidden, context = self.encoder(src)
  init_output = self.make_init_decoder_output(context)
 

OpenNMT/onmt/Translator.py

@@ -48,17 +48,17 @@ def buildTargetTokens(self, pred, src, attn):
 
  def translateBatch(self, batch):
  srcBatch, tgtBatch = batch
- batchSize = srcBatch.size(0)
+ batchSize = srcBatch.size(1)
  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(1), dim=1):
+ for srcBatch_t in srcBatch.split(1):
  encStates, context_t = self.model.encoder(srcBatch_t, hidden=encStates)
- batchPadIdx = srcBatch_t.data.squeeze(1).eq(onmt.Constants.PAD).nonzero()
+ batchPadIdx = srcBatch_t.data.squeeze(0).eq(onmt.Constants.PAD).nonzero()
  if batchPadIdx.nelement() > 0:
  batchPadIdx = batchPadIdx.squeeze(1)
  encStates[0].data.index_fill_(1, batchPadIdx, 0)
@@ -73,7 +73,7 @@ def translateBatch(self, batch):
 
  # 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)
+ padMask = srcBatch.data.eq(onmt.Constants.PAD).t()
  def applyContextMask(m):
  if isinstance(m, onmt.modules.GlobalAttention):
  m.applyMask(padMask)
@@ -88,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.transpose(0, 1), tgtBatch.transpose(0, 1)[1:].data):
+ tgtBatch[:-1], decStates, context, initOutput)
+ for dec_t, tgt_t in zip(decOut, tgtBatch[1:].data):
  gen_t = self.model.generator.forward(dec_t)
  tgt_t = tgt_t.unsqueeze(1)
  scores = gen_t.data.gather(1, tgt_t)
@@ -107,7 +107,7 @@ def applyContextMask(m):
 
  decOut = self.model.make_init_decoder_output(context)
 
- padMask = srcBatch.data.eq(onmt.Constants.PAD).unsqueeze(0).repeat(beamSize, 1, 1)
+ padMask = srcBatch.data.eq(onmt.Constants.PAD).t().unsqueeze(0).repeat(beamSize, 1, 1)
 
  batchIdx = list(range(batchSize))
  remainingSents = batchSize
@@ -120,9 +120,9 @@ def applyContextMask(m):
  if not b.done]).t().contiguous().view(1, -1)
 
  decOut, decStates, attn = self.model.decoder(
- Variable(input, volatile=True).transpose(0, 1), decStates, context, decOut)
+ Variable(input, volatile=True), decStates, context, decOut)
  # decOut: 1 x (beam*batch) x numWords
- decOut = decOut.transpose(0, 1).squeeze(0)
+ decOut = decOut.squeeze(0)
  out = self.model.generator.forward(decOut)
 
  # batch x beam x numWords
@@ -177,7 +177,7 @@ def updateActive(t):
  scores, ks = beam[b].sortBest()
 
  allScores += [scores[:n_best]]
- valid_attn = srcBatch.transpose(0, 1).data[:, b].ne(onmt.Constants.PAD).nonzero().squeeze(1)
+ valid_attn = srcBatch.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]
@@ -189,14 +189,13 @@ 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(0)):
+ for b in range(batch[0].size(1)):
  predBatch.append(
  [self.buildTargetTokens(pred[b][n], srcBatch[b], attn[b][n])
  for n in range(self.opt.n_best)]

OpenNMT/train.py

@@ -117,11 +117,11 @@ def NMTCriterion(vocabSize):
 def memoryEfficientLoss(outputs, targets, generator, crit, eval=False):
  # compute generations one piece at a time
  loss = 0
- outputs = Variable(outputs.data, requires_grad=(not eval), volatile=eval).contiguous()
+ outputs = Variable(outputs.data, requires_grad=(not eval), volatile=eval)
 
  batch_size = outputs.size(1)
  outputs_split = torch.split(outputs, opt.max_generator_batches)
- targets_split = torch.split(targets.contiguous(), opt.max_generator_batches)
+ targets_split = torch.split(targets, opt.max_generator_batches)
  for out_t, targ_t in zip(outputs_split, targets_split):
  out_t = out_t.view(-1, out_t.size(2))
  pred_t = generator(out_t)
@@ -140,9 +140,9 @@ def eval(model, criterion, data):
 
  model.eval()
  for i in range(len(data)):
- batch = [x.transpose(0, 1) for x in data[i]] # must be batch first for gather/scatter in DataParallel
+ batch = data[i]
  outputs = model(batch) # FIXME volatile
- targets = batch[1][:, 1:] # exclude <s> from targets
+ targets = batch[1][1:] # exclude <s> from targets
  loss, _ = memoryEfficientLoss(
  outputs, targets, model.generator, criterion, eval=True)
  total_loss += loss
@@ -172,11 +172,10 @@ def trainEpoch(epoch):
 
  batchIdx = batchOrder[i] if epoch >= opt.curriculum else i
  batch = trainData[batchIdx]
- batch = [x.transpose(0, 1) for x in batch] # must be batch first for gather/scatter in DataParallel
 
  model.zero_grad()
  outputs = model(batch)
- targets = batch[1][:, 1:] # exclude <s> from targets
+ targets = batch[1][1:] # exclude <s> from targets
  loss, gradOutput = memoryEfficientLoss(
  outputs, targets, model.generator, criterion)
 
@@ -209,7 +208,8 @@ def trainEpoch(epoch):
 
  # (1) train for one epoch on the training set
  train_loss = trainEpoch(epoch)
- print('Train perplexity: %g' % math.exp(min(train_loss, 100)))
+ train_ppl = math.exp(min(train_loss, 100))
+ print('Train perplexity: %g' % train_ppl)
 
  # (2) evaluate on the validation set
  valid_loss = eval(model, criterion, validData)
@@ -229,8 +229,7 @@ def trainEpoch(epoch):
  'optim': optim,
  }
  torch.save(checkpoint,
- '%s_e%d_%.2f.pt' % (opt.save_model, epoch, valid_ppl))
-
+ '%s_val%.2f_e%d_train%.2f.pt' % (opt.save_model, valid_ppl, epoch, train_ppl))
 
 def main():
 
@@ -258,11 +257,11 @@ def main():
  generator = nn.Sequential(
  nn.Linear(opt.rnn_size, dicts['tgt'].size()),
  nn.LogSoftmax())
- if len(opt.gpus) > 1:
- generator = nn.DataParallel(generator, device_ids=opt.gpus)
+# if len(opt.gpus) > 1:
+# generator = nn.DataParallel(generator, device_ids=opt.gpus)
  model = onmt.Models.NMTModel(encoder, decoder, generator)
  if len(opt.gpus) > 1:
- model = nn.DataParallel(model, device_ids=opt.gpus)
+ model = nn.DataParallel(model, device_ids=opt.gpus, dim=1)
  if opt.gpus:
  model.cuda()
  else: