multi-gpu via DataParallel

Author: bmccann

OpenNMT/onmt/Models.py

@@ -29,10 +29,9 @@ def __init__(self, opt, dicts):
             self.word_lut.weight.copy_(pretrained)
 
     def forward(self, input, hidden=None):
-        emb = self.word_lut(input)
-
+        batch_size = input.size(0) # batch first for multi-gpu compatibility
+        emb = self.word_lut(input).transpose(0, 1)
         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)
@@ -46,21 +45,21 @@ 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 = []
         for i in range(num_layers):
             layer = nn.LSTMCell(input_size, rnn_size)
             self.add_module('layer_%d' % i, layer)
-            self.layers += [layer]
             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):
+        for i in range(self.num_layers):
+            layer = getattr(self, 'layer_%d' % i)
             h_1_i, c_1_i = layer(input, (h_0[i], c_0[i]))
             input = h_1_i
-            if i != len(self.layers):
+            if i != self.num_layers:
                 input = self.dropout(input)
             h_1 += [h_1_i]
             c_1 += [c_1_i]
@@ -99,9 +98,9 @@ def __init__(self, opt, dicts):
 
 
     def forward(self, input, hidden, context, init_output):
-        emb = self.word_lut(input)
+        emb = self.word_lut(input).transpose(0, 1)
 
-        batch_size = input.size(1)
+        batch_size = input.size(0)
 
         h_size = (batch_size, self.hidden_size)
         output = Variable(emb.data.new(*h_size).zero_(), requires_grad=False)
@@ -122,7 +121,7 @@ def forward(self, input, hidden, context, init_output):
             outputs += [output]
 
         outputs = torch.stack(outputs)
-        return outputs, hidden, attn
+        return outputs.transpose(0, 1), hidden, attn
 
 
 class NMTModel(nn.Module):
@@ -154,7 +153,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/preprocess.py

@@ -1,7 +1,6 @@
 import onmt
 
 import argparse
-import os
 import torch
 
 parser = argparse.ArgumentParser(description='preprocess.lua')

OpenNMT/train.py

@@ -2,6 +2,7 @@
 import argparse
 import torch
 import torch.nn as nn
+from torch import cuda
 from torch.autograd import Variable
 import math
 import time
@@ -86,7 +87,7 @@
                     See README for specific formatting instructions.""")
 
 # GPU
-parser.add_argument('-cuda', action='store_true',
+parser.add_argument('-gpu', default=[], nargs='+', type=int,
                     help="Use CUDA")
 
 parser.add_argument('-log_interval', type=int, default=50,
@@ -95,11 +96,15 @@
 #                     help="Seed for random initialization")
 
 opt = parser.parse_args()
+opt.cuda = len(opt.gpu)
+
 print(opt)
 
 if torch.cuda.is_available() and not opt.cuda:
     print("WARNING: You have a CUDA device, so you should probably run with -cuda")
 
+if opt.cuda:
+    cuda.set_device(opt.gpu[0])
 
 def NMTCriterion(vocabSize):
     weight = torch.ones(vocabSize)
@@ -117,7 +122,7 @@ def memoryEfficientLoss(outputs, targets, generator, crit, eval=False):
 
     batch_size = outputs.size(1)
     outputs_split = torch.split(outputs, opt.max_generator_batches)
-    targets_split = torch.split(targets, opt.max_generator_batches)
+    targets_split = torch.split(targets.contiguous(), 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)
@@ -136,9 +141,9 @@ def eval(model, criterion, data):
 
     model.eval()
     for i in range(len(data)):
-        batch = data[i]
+        batch = [x.transpose(0, 1) for x in data[i]] # must be batch first for gather/scatter in DataParallel
         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
@@ -155,6 +160,7 @@ def trainModel(model, trainData, validData, dataset, optim):
     # define criterion of each GPU
     criterion = NMTCriterion(dataset['dicts']['tgt'].size())
 
+    start_time = time.time()
     def trainEpoch(epoch):
 
         # shuffle mini batch order
@@ -167,10 +173,11 @@ 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)
 
@@ -185,10 +192,11 @@ def trainEpoch(epoch):
             total_words += num_words
             report_words += num_words
             if i % opt.log_interval == 0 and i > 0:
-                print("Epoch %2d, %5d/%5d batches; perplexity: %6.2f; %3.0f tokens/s" %
+                print("Epoch %2d, %5d/%5d batches; perplexity: %6.2f; %3.0f tokens/s; %6.0f s elapsed" %
                       (epoch, i, len(trainData),
                       math.exp(report_loss / report_words),
-                      report_words/(time.time()-start)))
+                      report_words/(time.time()-start),
+                      time.time()-start_time))
 
                 report_loss = report_words = 0
                 start = time.time()
@@ -249,7 +257,15 @@ def main():
         generator = nn.Sequential(
             nn.Linear(opt.rnn_size, dicts['tgt'].size()),
             nn.LogSoftmax())
+        generator = nn.DataParallel(generator, device_ids=opt.gpu)
         model = onmt.Models.NMTModel(encoder, decoder, generator)
+        model = nn.DataParallel(model, device_ids=opt.gpu)
+        if opt.cuda:
+            model.cuda()
+        else:
+            model.cpu()
+
+        model.generator = generator
 
         for p in model.parameters():
             p.data.uniform_(-opt.param_init, opt.param_init)
@@ -263,14 +279,13 @@ def main():
         print('Loading from checkpoint at %s' % opt.train_from)
         checkpoint = torch.load(opt.train_from)
         model = checkpoint['model']
+        if opt.cuda:
+            model.cuda()
+        else:
+            model.cpu()
         optim = checkpoint['optim']
         opt.start_epoch = checkpoint['epoch'] + 1
 
-    if opt.cuda:
-        model.cuda()
-    else:
-        model.cpu()
-
     nParams = sum([p.nelement() for p in model.parameters()])
     print('* number of parameters: %d' % nParams)