Esempio n. 1
0
def beam_search(model: NMT, test_data_src: List[List[str]], beam_size: int,
                max_decoding_time_step: int) -> List[List[Hypothesis]]:
    """ Run beam search to construct hypotheses for a list of src-language sentences.
    @param model (NMT): NMT Model
    @param test_data_src (List[List[str]]): List of sentences (words) in source language, from test set.
    @param beam_size (int): beam_size (# of hypotheses to hold for a translation at every step)
    @param max_decoding_time_step (int): maximum sentence length that Beam search can produce
    @returns hypotheses (List[List[Hypothesis]]): List of Hypothesis translations for every source sentence.
    """
    was_training = model.training
    model.eval()

    hypotheses = []
    with torch.no_grad():
        for src_sent in tqdm(test_data_src, desc='Decoding', file=sys.stdout):
            example_hyps = model.beam_search(
                src_sent,
                beam_size=beam_size,
                max_decoding_time_step=max_decoding_time_step)

            hypotheses.append(example_hyps)

    if was_training: model.train(was_training)

    return hypotheses
def beam(args):
    # load model params
    print('load model from [%s]' % args.model_bin, file=sys.stderr)
    params = torch.load(args.model_bin,
                        map_location=lambda storage, loc: storage)
    vocab = params['vocab']
    opt = params['args']
    state_dict = params['state_dict']

    # build model
    model = NMT(opt, vocab)
    model.load_state_dict(state_dict)
    model.train()
    # model.eval()
    model = model.cuda()

    # loss function
    loss_fn = torch.nn.NLLLoss()

    # sampling
    print('begin beam searching')
    src_sent = ['we', 'have', 'told', 'that', '.']
    hyps = model.beam(src_sent)

    print('src_sent:', ' '.join(src_sent))
    for ids, hyp, dist in hyps:
        print('tgt_sent:', ' '.join(hyp))
        print('tgt_ids :', end=' ')
        for id in ids:
            print(id, end=', ')
        print()
        print('out_dist:', dist)

        var_ids = torch.autograd.Variable(torch.LongTensor(ids[1:]),
                                          requires_grad=False)
        loss = loss_fn(dist, var_ids)
        print('NLL loss =', loss)

    loss.backward()
Esempio n. 3
0
def train(mode, checkpoint_path):
    # Data
    data_train = IWSLT15EnViDataSet(en_path="../data/train-en-vi/train.en",
                                    vi_path="../data/train-en-vi/train.vi")
    data_loader = DataLoader(data_train,
                             batch_size=BATCH_SIZE,
                             shuffle=False,
                             drop_last=False)
    if mode == EN2VI:
        src_vocab_size, tgt_vocab_size = data_train.en_vocab_size, data_train.vi_vocab_size
    else:
        src_vocab_size, tgt_vocab_size = data_train.vi_vocab_size, data_train.en_vocab_size
    print("Loading data done!")

    # Model & Optimizer
    model = NMT(mode=mode,
                src_vocab_size=src_vocab_size,
                tgt_vocab_size=tgt_vocab_size)
    model.to(device)

    criterion = MaskedPaddingCrossEntropyLoss().to(device)
    optimizer = Adam(model.parameters())

    prev_epoch = 0
    if checkpoint_path.exists():  # Resume training
        model, optimizer, prev_epoch = load_checkpoint(model, optimizer,
                                                       checkpoint_path)
        print(f"Resume training from {prev_epoch} epochs!")
    else:
        model.apply(xavier_init_weights)
        print("Training from start!")

    model.train()
    for epoch in range(N_EPOCHS - prev_epoch):
        print(f"\nEpoch: {epoch+prev_epoch+1}")

        for b, (en_tokens, en_valid_len, vi_tokens,
                vi_valid_len) in enumerate(data_loader):
            en_tokens, vi_tokens = en_tokens.to(device), vi_tokens.to(device)
            en_valid_len, vi_valid_len = en_valid_len.to(
                device), vi_valid_len.to(device)

            en_padding_masks = mask_padding(en_tokens, en_valid_len, device)
            vi_padding_masks = mask_padding(vi_tokens, vi_valid_len, device)

            if mode == EN2VI:
                src, tgt = en_tokens, vi_tokens
                tgt_valid_len = vi_valid_len
                src_masks, tgt_masks = en_padding_masks, vi_padding_masks
            else:
                src, tgt = vi_tokens, en_tokens
                tgt_valid_len = en_valid_len
                src_masks, tgt_masks = vi_padding_masks, en_padding_masks

            optimizer.zero_grad()

            # Encoder's forward pass:
            encoder_state = model.encoder(src, src_masks)
            # Decoder's forward pass
            decoder_X = torch.tensor([[DEFAULT_SOS_INDEX] * tgt.shape[0]],
                                     device=device).reshape(-1, 1)
            decoder_state = encoder_state

            loss = torch.tensor(0, device=device, dtype=torch.float)
            for i in range(1, tgt.shape[1]):
                decoder_state, logit_pred = model.decoder(
                    decoder_X, decoder_state)
                loss += criterion(pred=logit_pred[:, 0, :],
                                  label=tgt[:, i],
                                  device=device).sum()
                # Teacher forcing
                decoder_X = tgt[:, i].reshape(-1, 1)

            loss.backward()
            clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

            if b % 50 == 0:
                seq_loss = loss / (MAX_LENGTH - 1)
                print(f"\tBatch {b}; Loss: {seq_loss:.2f}; "
                      f"Mean Token Loss: {seq_loss/tgt_valid_len.sum():.4f}")

            ## Free up GPU memory
            del src, tgt, en_valid_len, vi_valid_len, decoder_state, logit_pred, loss
            torch.cuda.empty_cache()

        save_checkpoint(mode, src_vocab_size, tgt_vocab_size, model, optimizer,
                        data_train.tokenizer_en, data_train.tokenizer_vi,
                        prev_epoch + epoch + 1, checkpoint_path)

        for en in ens:
            vi = translate_en2vi(en_sentence=en,
                                 length=MAX_LENGTH,
                                 model=model,
                                 tokenizer_en=data_train.tokenizer_en,
                                 tokenizer_vi=data_train.tokenizer_vi,
                                 device=device)
            print("en:", en, "=> vi:", vi)
Esempio n. 4
0
    timer = Timer(epoch_size)
    init_output_log(save_dir)

    print('model load done.')
    # get model

    train_len = len(train_x)
    dev_len = len(dev_x)
    train_x, train_y, train_mask = data_padding(train_x, train_y)
    dev_x, dev_y, dev_mask = data_padding(dev_x, dev_y)

    print('training start.')

    best_loss = 998244353.0
    for epoch in range(epoch_size):
        model.train()
        sum_loss = 0
        batch_num = math.ceil(train_len / batch_size)
        for step in range(batch_num):
            inputs = batch_iter(train_x, step, batch_size)
            labels = batch_iter(train_y, step, batch_size)
            masks = batch_iter(train_mask, step, batch_size)
            inputs = torch.LongTensor(inputs).to(device)
            labels = torch.LongTensor(labels).to(device)
            masks = torch.ByteTensor(masks).to(device)

            optimizer.zero_grad()
            outputs = model(inputs, labels)
            Loss = Loss_fn(outputs, labels, masks)
            Loss.backward()
            optimizer.step()
Esempio n. 5
0
class Trainer:
    """
    训练类,使用训练集训练模型

    Args:
        _hparams (NameSpace): 人为设定的超参数,默认值见config.py,也可以在命令行指定。
    """

    def __init__(self, _hparams):
        self.hparams = _hparams
        set_seed(_hparams.fixed_seed)
        self.train_loader = get_dataloader(_hparams.train_src_path, _hparams.train_dst_path,
                                           _hparams.batch_size, _hparams.num_workers)
        self.src_vocab, self.dst_vocab = load_vocab(_hparams.train_src_pkl, _hparams.train_dst_pkl)
        self.device = torch.device(_hparams.device)
        self.model = NMT(_hparams.embed_size, _hparams.hidden_size,
                         self.src_vocab, self.dst_vocab, self.device,
                         _hparams.dropout_rate).to(self.device)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=_hparams.lr)

    def train(self):
        print('*' * 20, 'train', '*' * 20)
        hist_valid_scores = []
        patience = 0
        num_trial = 0

        for epoch in range(int(self.hparams.max_epochs)):
            self.model.train()

            epoch_loss_val = 0
            epoch_steps = len(self.train_loader)
            for step, data_pairs in tqdm(enumerate(self.train_loader)):
                sents = [(dp.src, dp.dst) for dp in data_pairs]
                src_sents, tgt_sents = zip(*sents)

                self.optimizer.zero_grad()

                batch_size = len(src_sents)
                example_losses = -self.model(src_sents, tgt_sents)
                batch_loss = example_losses.sum()
                train_loss = batch_loss / batch_size
                epoch_loss_val += train_loss.item()
                train_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.hparams.clip_gradient)
                self.optimizer.step()

            epoch_loss_val /= epoch_steps
            print('epoch: {}, epoch_loss_val: {}'.format(epoch, epoch_loss_val))

            # perform validation
            if epoch % self.hparams.valid_niter == 0:
                print('*' * 20, 'validate', '*' * 20)
                dev_ppl = evaluate_ppl(self.model, self.hparams.val_src_path, self.hparams.val_dst_path,
                                       self.hparams.batch_val_size, self.hparams.num_workers)
                valid_metric = -dev_ppl

                is_better = len(hist_valid_scores) == 0 or valid_metric > max(hist_valid_scores)
                hist_valid_scores.append(valid_metric)

                if is_better:
                    patience = 0
                    print('save currently the best model to {}'.format(self.hparams.model_save_path))
                    self.model.save(self.hparams.model_save_path)
                    torch.save(self.optimizer.state_dict(), self.hparams.optimizer_save_path)
                elif patience < self.hparams.patience:
                    patience += 1
                    print('hit patience %d' % patience)

                    if patience == self.hparams.patience:
                        num_trial += 1
                        print('hit #{} trial'.format(num_trial))
                        if num_trial == self.hparams.max_num_trial:
                            print('early stop!')
                            exit(0)

                        # 兼容设计,考虑Adam不需要人工调整lr,而其他优化器需要
                        if hasattr(self.optimizer, 'param_group'):
                            # decay lr, and restore from previously best checkpoint
                            lr = self.optimizer.param_groups[0]['lr'] * self.hparams.lr_decay
                            print('load previously best model and decay learning rate to %f' % lr)

                            params = torch.load(self.hparams.model_save_path, map_location=lambda storage, loc: storage)
                            self.model.load_state_dict(params['state_dict'])
                            self.model = self.model.to(self.device)

                            print('restore parameters of the optimizers')
                            self.optimizer.load_state_dict(torch.load(self.hparams.optimizer_save_path))

                            # set new lr
                            for param_group in self.optimizer.param_groups:
                                param_group['lr'] = lr

                        # reset patience
                        patience = 0
                print('*' * 20, 'end validate', '*' * 20)
        print('*' * 20, 'end train', '*' * 20)
Esempio n. 6
0
def train(args: Dict):
    train_data_src = read_corpus(args['--train-src'], source='src')
    train_data_tgt = read_corpus(args['--train-tgt'], source='tgt')

    dev_data_src = read_corpus(args['--dev-src'], source='src')
    dev_data_tgt = read_corpus(args['--dev-tgt'], source='tgt')

    # [(src_0, tgt_0), (src_1, tgt_1), ..., ]
    train_data = list(zip(train_data_src, train_data_tgt))
    dev_data = list(zip(dev_data_src, dev_data_tgt))

    train_batch_size = int(args['--batch-size'])
    clip_grad = float(args['--clip-grad'])
    valid_niter = int(args['--valid-niter'])
    log_every = int(args['--log-every'])
    model_save_path = args['--save-to']

    # vocab = Vocab.load(args['--vocab'])
    vocab = Vocab.build(train_data_src, train_data_tgt,
                        int(args['--vocab-size']), 1)

    model = NMT(embed_size=int(args['--embed-size']),
                hidden_size=int(args['--hidden-size']),
                dropout_rate=float(args['--dropout']),
                vocab=vocab)
    model.train()
    print(model)

    uniform_init = float(args['--uniform-init'])
    if np.abs(uniform_init) > 0.:
        print('uniformly initialize parameters [-%f, +%f]' %
              (uniform_init, uniform_init),
              file=sys.stderr)
        for p in model.parameters():
            p.data.uniform_(-uniform_init, uniform_init)

    # vocab_mask = torch.ones(len(vocab.tgt))
    # vocab_mask[vocab.tgt['<pad>']] = 0

    device = torch.device("cuda:0" if args['--cuda'] else "cpu")
    print('use device: %s' % device, file=sys.stderr)

    model = model.to(device)
    model.save(model_save_path)

    optimizer = torch.optim.Adam(model.parameters(), lr=float(args['--lr']))

    num_trial = 0
    train_iter = patience = cum_loss = report_loss = cum_tgt_words = report_tgt_words = 0
    cum_examples = report_examples = epoch = valid_num = 0
    hist_valid_scores = []
    train_time = begin_time = time.time()

    print('begin Maximum Likelihood training')

    while True:
        epoch += 1

        for src_sents, tgt_sents in batch_iter(train_data,
                                               batch_size=train_batch_size,
                                               shuffle=True):
            train_iter += 1

            optimizer.zero_grad()

            batch_size = len(src_sents)

            #################### forward pass and compute loss #########################
            # example_losses = -model(src_sents, tgt_sents) # (batch_size,)
            example_losses = model(src_sents, tgt_sents)  # [batch_size,]
            batch_loss = example_losses.sum()
            loss = batch_loss / batch_size

            #################### backward pass to compute gradients ####################
            loss.backward()

            # clip gradient
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),
                                                       clip_grad)

            #################### update model parameters ###############################
            optimizer.step()

            #################### do some statistics ####################################
            batch_losses_val = batch_loss.item()
            report_loss += batch_losses_val
            cum_loss += batch_losses_val

            tgt_words_num_to_predict = sum(
                len(s[1:]) for s in tgt_sents)  # omitting leading `<s>`
            report_tgt_words += tgt_words_num_to_predict
            cum_tgt_words += tgt_words_num_to_predict
            report_examples += batch_size
            cum_examples += batch_size

            #################### print log #############################################
            if train_iter % log_every == 0:
                print(
                    'epoch %d, iter %d, avg. loss %.2f, avg. ppl %.2f '
                    'cum. examples %d, speed %.2f words/sec, time elapsed %.2f sec'
                    % (
                        epoch,
                        train_iter,
                        report_loss / report_examples,
                        #  math.exp(report_loss / report_tgt_words),
                        (report_loss / report_tgt_words),
                        cum_examples,
                        report_tgt_words / (time.time() - train_time),
                        time.time() - begin_time),
                    file=sys.stderr)

                train_time = time.time()
                report_loss = report_tgt_words = report_examples = 0.

            ##################### perform validation ##################################
            if train_iter % valid_niter == 0:
                print(
                    'epoch %d, iter %d, cum. loss %.2f, cum. ppl %.2f cum. examples %d'
                    % (epoch, train_iter, cum_loss / cum_examples,
                       np.exp(cum_loss / cum_tgt_words), cum_examples),
                    file=sys.stderr)

                cum_loss = cum_examples = cum_tgt_words = 0.
                valid_num += 1

                print('begin validation ...', file=sys.stderr)

                # compute dev. ppl and bleu
                dev_ppl = evaluate_ppl(
                    model, dev_data,
                    batch_size=128)  # dev batch size can be a bit larger
                valid_metric = -dev_ppl

                print('validation: iter %d, dev. ppl %f' %
                      (train_iter, dev_ppl),
                      file=sys.stderr)

                is_better = len(hist_valid_scores
                                ) == 0 or valid_metric > max(hist_valid_scores)
                hist_valid_scores.append(valid_metric)

                # hypotheses = beam_search(model, dev_data_src,
                #                          beam_size=4,
                #                          max_decoding_time_step=10)

                if is_better:
                    patience = 0
                    print('save currently the best model to [%s]' %
                          model_save_path,
                          file=sys.stderr)
                    model.save(model_save_path)

                    # also save the optimizers' state
                    torch.save(optimizer.state_dict(),
                               model_save_path + '.optim')
                elif patience < int(args['--patience']):
                    patience += 1
                    print('hit patience %d' % patience, file=sys.stderr)

                    if patience == int(args['--patience']):
                        num_trial += 1
                        print('hit #%d trial' % num_trial, file=sys.stderr)
                        if num_trial == int(args['--max-num-trial']):
                            print('early stop!', file=sys.stderr)
                            exit(0)

                        # decay lr, and restore from previously best checkpoint
                        lr = optimizer.param_groups[0]['lr'] * float(
                            args['--lr-decay'])
                        print(
                            'load previously best model and decay learning rate to %f'
                            % lr,
                            file=sys.stderr)

                        # load model
                        params = torch.load(
                            model_save_path,
                            map_location=lambda storage, loc: storage)
                        model.load_state_dict(params['state_dict'])
                        model = model.to(device)

                        print('restore parameters of the optimizers',
                              file=sys.stderr)
                        optimizer.load_state_dict(
                            torch.load(model_save_path + '.optim'))

                        # set new lr
                        for param_group in optimizer.param_groups:
                            param_group['lr'] = lr

                        # reset patience
                        patience = 0

                if epoch == int(args['--max-epoch']):
                    print('reached maximum number of epochs!', file=sys.stderr)
                    exit(0)
Esempio n. 7
0
def main(options):

  use_cuda = (len(options.gpuid) >= 1)
  if options.gpuid:
    cuda.set_device(options.gpuid[0])

  src_train, src_dev, src_test, src_vocab = torch.load(open(options.data_file + "." + options.src_lang, 'rb'))
  trg_train, trg_dev, trg_test, trg_vocab = torch.load(open(options.data_file + "." + options.trg_lang, 'rb'))

  batched_train_src, batched_train_src_mask, sort_index = utils.tensor.advanced_batchize(src_train, options.batch_size, src_vocab.stoi["<blank>"])
  batched_train_trg, batched_train_trg_mask = utils.tensor.advanced_batchize_no_sort(trg_train, options.batch_size, trg_vocab.stoi["<blank>"], sort_index)
  batched_dev_src, batched_dev_src_mask, sort_index = utils.tensor.advanced_batchize(src_dev, options.batch_size, src_vocab.stoi["<blank>"])
  batched_dev_trg, batched_dev_trg_mask = utils.tensor.advanced_batchize_no_sort(trg_dev, options.batch_size, trg_vocab.stoi["<blank>"], sort_index)

  trg_vocab_size = len(trg_vocab)
  src_vocab_size = len(src_vocab)
  word_emb_size = 300
  hidden_size = 1024

  nmt = NMT(src_vocab_size, trg_vocab_size, word_emb_size, hidden_size,
            src_vocab, trg_vocab, attn_model = "general", use_cuda = True)

  if use_cuda > 0:
    nmt.cuda()
    if options.distributed:
      nmt = torch.nn.DataParallel(nmt)
  else:
    nmt.cpu()

  criterion = torch.nn.NLLLoss()

  # Configure optimization
  lr = options.learning_rate
  optimizer = eval("torch.optim." + options.optimizer)(nmt.parameters(), lr)

  
  # main training loop
  last_dev_avg_loss = float("inf")
  for epoch_i in range(options.epochs):
    logging.info("At {0}-th epoch.".format(epoch_i))

    # Set training mode
    nmt.train()

    # srange generates a lazy sequence of shuffled range
    for i, batch_i in enumerate(utils.rand.srange(len(batched_train_src))):
      train_src_batch = Variable(batched_train_src[batch_i])  # of size (src_seq_len, batch_size)
      train_trg_batch = Variable(batched_train_trg[batch_i])  # of size (src_seq_len, batch_size)
      train_src_mask = Variable(batched_train_src_mask[batch_i])
      train_trg_mask = Variable(batched_train_trg_mask[batch_i])
      if use_cuda:
        train_src_batch = train_src_batch.cuda()
        train_trg_batch = train_trg_batch.cuda()
        train_src_mask = train_src_mask.cuda()
        train_trg_mask = train_trg_mask.cuda()

      sys_out_batch = nmt(train_src_batch, train_trg_batch, True)

      del train_src_batch

      train_trg_mask = train_trg_mask.view(-1)
      train_trg_batch = train_trg_batch.view(-1)
      train_trg_batch = train_trg_batch.masked_select(train_trg_mask)
      train_trg_mask = train_trg_mask.unsqueeze(1).expand(len(train_trg_mask), trg_vocab_size)
      sys_out_batch = sys_out_batch.view(-1, trg_vocab_size)
      sys_out_batch = sys_out_batch.masked_select(train_trg_mask).view(-1, trg_vocab_size)
      loss = criterion(sys_out_batch, train_trg_batch)
      logging.debug("loss at batch {0}: {1}".format(i, loss.data[0]))
      
      optimizer.zero_grad()
      loss.backward()
      # # gradient clipping
      torch.nn.utils.clip_grad_norm(nmt.parameters(), 5.0)
      optimizer.step()

    # validation -- this is a crude esitmation because there might be some paddings at the end
    dev_loss = 0.0

    # Set validation mode
    nmt.eval()

    for batch_i in range(len(batched_dev_src)):
      dev_src_batch = Variable(batched_dev_src[batch_i], volatile=True)
      dev_trg_batch = Variable(batched_dev_trg[batch_i], volatile=True)
      dev_src_mask = Variable(batched_dev_src_mask[batch_i], volatile=True)
      dev_trg_mask = Variable(batched_dev_trg_mask[batch_i], volatile=True)
      if use_cuda:
        dev_src_batch = dev_src_batch.cuda()
        dev_trg_batch = dev_trg_batch.cuda()
        dev_src_mask = dev_src_mask.cuda()
        dev_trg_mask = dev_trg_mask.cuda()

      sys_out_batch = nmt(dev_src_batch, dev_trg_batch, False)

      dev_trg_mask = dev_trg_mask.view(-1)
      dev_trg_batch = dev_trg_batch.view(-1)
      dev_trg_batch = dev_trg_batch.masked_select(dev_trg_mask)
      dev_trg_mask = dev_trg_mask.unsqueeze(1).expand(len(dev_trg_mask), trg_vocab_size)
      sys_out_batch = sys_out_batch.view(-1, trg_vocab_size)
      sys_out_batch = sys_out_batch.masked_select(dev_trg_mask).view(-1, trg_vocab_size)
      loss = criterion(sys_out_batch, dev_trg_batch)
      logging.debug("dev loss at batch {0}: {1}".format(batch_i, loss.data[0]))
      dev_loss += loss
    dev_avg_loss = dev_loss / len(batched_dev_src)
    logging.info("Average loss value per instance is {0} at the end of epoch {1}".format(dev_avg_loss.data[0], epoch_i))

    # if (last_dev_avg_loss - dev_avg_loss).data[0] < options.estop:
    #   logging.info("Early stopping triggered with threshold {0} (previous dev loss: {1}, current: {2})".format(epoch_i, last_dev_avg_loss.data[0], dev_avg_loss.data[0]))
    #   break
    torch.save(nmt, open(options.model_file + ".nll_{0:.2f}.epoch_{1}".format(dev_avg_loss.data[0], epoch_i), 'wb'), pickle_module=dill)
    last_dev_avg_loss = dev_avg_loss
Esempio n. 8
0
def train(args: Dict):
    """ Train the NMT Model.
    @param args (Dict): args from cmd line
    """
    do_bleu = '--ignore-test-bleu' not in args or not args['--ignore-test-bleu']
    train_data_src = read_corpus(args['--train-src'],
                                 source='src',
                                 dev_mode=dev_mode)
    train_data_tgt = read_corpus(args['--train-tgt'],
                                 source='tgt',
                                 dev_mode=dev_mode)

    dev_data_src = read_corpus(args['--dev-src'],
                               source='src',
                               dev_mode=dev_mode)
    dev_data_tgt = read_corpus(args['--dev-tgt'],
                               source='tgt',
                               dev_mode=dev_mode)

    if do_bleu:
        test_data_src = read_corpus(args['--test-src'],
                                    source='src',
                                    dev_mode=dev_mode)
        test_data_tgt = read_corpus(args['--test-tgt'],
                                    source='tgt',
                                    dev_mode=dev_mode)

    train_data = list(zip(train_data_src, train_data_tgt))
    dev_data = list(zip(dev_data_src, dev_data_tgt))

    max_tokens_in_sentence = int(args['--max-decoding-time-step'])
    train_data = clean_data(train_data, max_tokens_in_sentence)
    dev_data = clean_data(dev_data, max_tokens_in_sentence)

    train_batch_size = int(args['--batch-size'])
    dev_batch_size = 128
    clip_grad = float(args['--clip-grad'])
    valid_niter = int(args['--valid-niter'])
    bleu_niter = int(args['--bleu-niter'])
    log_every = int(args['--log-every'])
    model_save_path = args['--save-to']

    vocab = Vocab.load(args['--vocab'], args['--word_freq'])

    model = NMT(embed_size=int(args['--embed-size']),
                hidden_size=int(args['--hidden-size']),
                dropout_rate=float(args['--dropout']),
                vocab=vocab)
    writer = SummaryWriter()

    # model = TransformerNMT(vocab, num_hidden_layers=3)

    model.train()

    uniform_init = float(args['--uniform-init'])
    if np.abs(uniform_init) > 0.:
        print('uniformly initialize parameters [-%f, +%f]' %
              (uniform_init, uniform_init),
              file=sys.stderr)
        for p in model.parameters():
            if p.dim() > 1:
                torch.nn.init.xavier_uniform_(p)
            else:
                p.data.uniform_(-uniform_init, uniform_init)

    vocab_mask = torch.ones(len(vocab.tgt))
    vocab_mask[vocab.tgt['<pad>']] = 0

    device = torch.device("cuda:0" if args['--cuda'] else "cpu")
    print('use device: %s' % device, file=sys.stderr)

    model = model.to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=float(args['--lr']))

    num_trial = 0
    train_iter = patience = cum_loss = report_loss = cum_tgt_words = report_tgt_words = 0
    cum_examples = report_examples = epoch = valid_num = 0
    hist_valid_scores = []
    train_time = begin_time = time.time()

    print("Sorting dataset based on difficulty...")
    dataset = (train_data, dev_data)
    ordered_dataset = load_order(args['--order-name'], dataset, vocab)
    # TODO: order = balance_order(order, dataset)
    (train_data, dev_data) = ordered_dataset

    visualize_scoring_examples = False
    if visualize_scoring_examples:
        visualize_scoring(ordered_dataset, vocab)

    n_iters = math.ceil(len(train_data) / train_batch_size)
    print("n_iters per epoch is {}: ({} / {})".format(n_iters, len(train_data),
                                                      train_batch_size))
    max_epoch = int(args['--max-epoch'])
    max_iters = max_epoch * n_iters

    print('begin Maximum Likelihood training')
    print('Using order function: {}'.format(args['--order-name']))
    print('Using pacing function: {}'.format(args['--pacing-name']))
    while True:
        epoch += 1
        for _ in range(n_iters):
            # Get pacing data according to train_iter
            current_train_data, current_dev_data = pacing_data(
                train_data,
                dev_data,
                time=train_iter,
                warmup_iters=int(args["--warmup-iters"]),
                method=args['--pacing-name'],
                tb=writer)

            # Uniformly sample batches from the paced dataset
            src_sents, tgt_sents = get_pacing_batch(
                current_train_data, batch_size=train_batch_size, shuffle=True)

            train_iter += 1

            # ERROR START
            optimizer.zero_grad()

            batch_size = len(src_sents)

            example_losses = -model(src_sents, tgt_sents)  # (batch_size,)
            batch_loss = example_losses.sum()
            loss = batch_loss / batch_size

            loss.backward()
            # clip gradient
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(),
                                                       clip_grad)

            optimizer.step()

            batch_losses_val: int = batch_loss.item()
            report_loss += batch_losses_val
            cum_loss += batch_losses_val

            tgt_words_num_to_predict = sum(
                len(s[1:]) for s in tgt_sents)  # omitting leading `<s>`
            report_tgt_words += tgt_words_num_to_predict
            cum_tgt_words += tgt_words_num_to_predict
            report_examples += batch_size
            cum_examples += batch_size

            if train_iter % log_every == 0:
                print(
                    'epoch %d, iter %d, avg. loss %.2f, avg. ppl %.2f '
                    'cum. examples %d, speed %.2f words/sec, time elapsed %.2f sec'
                    % (epoch, train_iter, report_loss / report_examples,
                       math.exp(report_loss / report_tgt_words), cum_examples,
                       report_tgt_words /
                       (time.time() - train_time), time.time() - begin_time),
                    file=sys.stderr)
                writer.add_scalar('Loss/train', report_loss / report_examples,
                                  train_iter)
                writer.add_scalar('ppl/train',
                                  math.exp(report_loss / report_tgt_words),
                                  train_iter)
                train_time = time.time()
                report_loss = report_tgt_words = report_examples = 0.

            # evaluate BLEU
            if train_iter % bleu_niter == 0 and do_bleu:
                bleu = decode_with_params(
                    model, test_data_src, test_data_tgt,
                    int(args['--beam-size']),
                    int(args['--max-decoding-time-step']))
                writer.add_scalar('bleu/test', bleu, train_iter)

            # perform validation
            if train_iter % valid_niter == 0:
                print(
                    'epoch %d, iter %d, cum. loss %.2f, cum. ppl %.2f cum. examples %d'
                    % (epoch, train_iter, cum_loss / cum_examples,
                       np.exp(cum_loss / cum_tgt_words), cum_examples),
                    file=sys.stderr)

                cum_loss = cum_examples = cum_tgt_words = 0.
                valid_num += 1

                print('begin validation ...', file=sys.stderr)

                # compute dev. ppl and bleu
                # dev batch size can be a bit larger
                dev_ppl = evaluate_ppl(model,
                                       current_dev_data,
                                       batch_size=dev_batch_size)
                valid_metric = -dev_ppl
                writer.add_scalar('ppl/valid', dev_ppl, train_iter)
                cum_loss = cum_examples = cum_tgt_words = 0.
                valid_num += 1

                print('validation: iter %d, dev. ppl %f' %
                      (train_iter, dev_ppl),
                      file=sys.stderr)

                is_better = len(hist_valid_scores
                                ) == 0 or valid_metric > max(hist_valid_scores)
                hist_valid_scores.append(valid_metric)

                if is_better:
                    patience = 0
                    print('save currently the best model to [%s]' %
                          model_save_path,
                          file=sys.stderr)
                    model.save(model_save_path)

                    # also save the optimizers' state
                    torch.save(optimizer.state_dict(),
                               model_save_path + '.optim')
                elif patience < int(args['--patience']):
                    patience += 1
                    print('hit patience %d' % patience, file=sys.stderr)

                    if patience == int(args['--patience']):
                        num_trial += 1
                        print('hit #%d trial' % num_trial, file=sys.stderr)
                        if num_trial == int(args['--max-num-trial']):
                            print('early stop!', file=sys.stderr)
                            exit(0)

                        # decay lr, and restore from previously best checkpoint
                        lr = optimizer.param_groups[0]['lr'] * \
                            float(args['--lr-decay'])
                        print(
                            'load previously best model and decay learning rate to %f'
                            % lr,
                            file=sys.stderr)

                        # load model
                        params = torch.load(
                            model_save_path,
                            map_location=lambda storage, loc: storage)
                        model.load_state_dict(params['state_dict'])
                        model = model.to(device)

                        print('restore parameters of the optimizers',
                              file=sys.stderr)
                        optimizer.load_state_dict(
                            torch.load(model_save_path + '.optim'))

                        # set new lr
                        for param_group in optimizer.param_groups:
                            param_group['lr'] = lr

                        # reset patience
                        patience = 0

                if epoch >= int(args['--max-epoch']):
                    print('reached maximum number of epochs!', file=sys.stderr)
                    exit(0)