state_dicts for translation and optimizer

Author: bmccann

OpenNMT/onmt/Models.py

@@ -149,7 +149,7 @@ def forward(self, input):
  self._fix_enc_hidden(enc_hidden[1]))
 
  out, dec_hidden, _attn = self.decoder(tgt, enc_hidden, context, init_output)
- if hasattr(self, 'generate') and self.generate:
+ if hasattr(self, 'generator') and self.generate:
  out = self.generator(out)
 
  return out

OpenNMT/onmt/Translator.py

@@ -1,4 +1,5 @@
 import onmt
+import torch.nn as nn
 import torch
 from torch.autograd import Variable
 
@@ -9,17 +10,34 @@ def __init__(self, opt):
  self.tt = torch.cuda if opt.cuda else torch
 
  checkpoint = torch.load(opt.model)
- self.model = checkpoint['model']
 
- self.model.eval()
+ model_opt = checkpoint['opt']
+ self.src_dict = checkpoint['dicts']['src']
+ self.tgt_dict = checkpoint['dicts']['tgt']
+
+ encoder = onmt.Models.Encoder(model_opt, self.src_dict)
+ decoder = onmt.Models.Decoder(model_opt, self.tgt_dict)
+ model = onmt.Models.NMTModel(encoder, decoder)
+
+ generator = nn.Sequential(
+ nn.Linear(model_opt.rnn_size, self.tgt_dict.size()),
+ nn.LogSoftmax())
+
+ model.load_state_dict(checkpoint['model'])
+ generator.load_state_dict(checkpoint['generator'])
 
  if opt.cuda:
- self.model.cuda()
+ model.cuda()
+ generator.cuda()
  else:
- self.model.cpu()
+ model.cpu()
+ generator.cpu()
+
+ model.generator = generator
+
+ self.model = model
+ self.model.eval()
 
- self.src_dict = checkpoint['dicts']['src']
- self.tgt_dict = checkpoint['dicts']['tgt']
 
  def buildData(self, srcBatch, goldBatch):
  srcData = [self.src_dict.convertToIdx(b,

OpenNMT/train.py

@@ -100,7 +100,7 @@
 
 # GPU
 parser.add_argument('-gpus', default=[], nargs='+', type=int,
- help="Use CUDA")
+ help="Use CUDA on the listed devices.")
 
 parser.add_argument('-log_interval', type=int, default=50,
  help="Print stats at this interval.")
@@ -255,7 +255,8 @@ def trainEpoch(epoch):
  'dicts': dataset['dicts'],
  'opt': opt,
  'epoch': epoch,
- 'optim': optim,
+ 'optimizer': optim.optimizer.state_dict(),
+ 'last_ppl': optim.last_ppl,
  }
  torch.save(checkpoint,
  '%s_acc_%.2f_ppl_%.2f_e%d.pt' % (opt.save_model, 100*valid_acc, valid_ppl, epoch))
@@ -299,12 +300,14 @@ def main():
  print('Loading model from checkpoint at %s' % opt.train_from)
  model.load_state_dict(checkpoint['model'])
  generator.load_state_dict(checkpoint['generator'])
- optim = checkpoint['optim']
  opt.start_epoch = checkpoint['epoch'] + 1
 
  if len(opt.gpus) >= 1:
  model.cuda()
  generator.cuda()
+ else:
+ model.cpu()
+ generator.cpu()
 
  if len(opt.gpus) > 1:
  model = nn.DataParallel(model, device_ids=opt.gpus, dim=1)
@@ -316,11 +319,15 @@ def main():
  for p in model.parameters():
  p.data.uniform_(-opt.param_init, opt.param_init)
 
- optim = onmt.Optim(
- model.parameters(), opt.optim, opt.learning_rate, opt.max_grad_norm,
- lr_decay=opt.learning_rate_decay,
- start_decay_at=opt.start_decay_at
- )
+ optim = onmt.Optim(
+ model.parameters(), opt.optim, opt.learning_rate, opt.max_grad_norm,
+ lr_decay=opt.learning_rate_decay,
+ start_decay_at=opt.start_decay_at
+ )
+
+ if opt.train_from:
+ optim.last_ppl = checkpoint['last_ppl']
+ optim.optimizer.load_state_dict(checkpoint['optimizer'])
 
  nParams = sum([p.nelement() for p in model.parameters()])
  print('* number of parameters: %d' % nParams)