# Copyright (c) 2018, salesforce.com, inc.
# All rights reserved.
# Licensed under the BSD 3-Clause license.
# For full license text, see the LICENSE file in the repo root
# or https://opensource.org/licenses/BSD-3-Clause
import torch
import numpy as np
from torchtext import data
from torchtext import datasets
from torch.nn import functional as F
from torch.autograd import Variable
import revtok
import logging
import random
import string
import traceback
import math
import uuid
import argparse
import os
import copy
import time
from tqdm import tqdm, trange
from model import Transformer, FastTransformer, INF, TINY, softmax
from utils import NormalField, NormalTranslationDataset, TripleTranslationDataset, ParallelDataset
from utils import Metrics, Best, computeGLEU, computeBLEU, Cache, Batch, masked_sort, unsorted, computeGroupBLEU
from time import gmtime, strftime
import sys
from traceback import extract_tb
from code import interact
def interactive_exception(e_class, e_value, tb):
sys.__excepthook__(e_class, e_value, tb)
tb_stack = extract_tb(tb)
locals_stack = []
while tb is not None:
locals_stack.append(tb.tb_frame.f_locals)
tb = tb.tb_next
while len(tb_stack) > 0:
frame = tb_stack.pop()
ls = locals_stack.pop()
print('\nInterpreter at file "{}", line {}, in {}:'.format(
frame.filename, frame.lineno, frame.name))
print(' {}'.format(frame.line.strip()))
interact(local=ls)
#sys.excepthook = interactive_exception
# check dirs
for d in ['models', 'runs', 'logs']:
if not os.path.exists('./{}'.format(d)):
os.mkdir('./{}'.format(d))
# params
parser = argparse.ArgumentParser(description='Train a Transformer model.')
# data
parser.add_argument('--data_prefix', type=str, default='../data/')
parser.add_argument('--dataset', type=str, default='iwslt', help='"flickr" or "iwslt"')
parser.add_argument('--language', type=str, default='ende', help='a combination of two language markers to show the language pair.')
parser.add_argument('--load_vocab', action='store_true', help='load a pre-computed vocabulary')
parser.add_argument('--load_dataset', action='store_true', help='load a pre-processed dataset')
parser.add_argument('--use_revtok', action='store_true', help='use reversible tokenization')
parser.add_argument('--level', type=str, default='subword', help='for BPE, we must preprocess the dataset')
parser.add_argument('--good_course', action='store_true', help='use beam-search output for distillation')
parser.add_argument('--test_set', type=str, default=None, help='which test set to use')
parser.add_argument('--max_len', type=int, default=None, help='limit the train set sentences to this many tokens')
parser.add_argument('--remove_eos', action='store_true', help='possibly remove <eos> tokens for FastTransformer')
# model basic
parser.add_argument('--prefix', type=str, default='', help='prefix to denote the model, nothing or [time]')
parser.add_argument('--params', type=str, default='james-iwslt', help='pamarater sets: james-iwslt, t2t-base, etc')
parser.add_argument('--fast', dest='model', action='store_const', const=FastTransformer,
default=Transformer, help='use a single self-attn stack')
# model variants
parser.add_argument('--local', dest='windows', action='store_const', const=[1, 3, 5, 7, -1],
default=None, help='use local attention')
parser.add_argument('--causal', action='store_true', help='use causal attention')
parser.add_argument('--positional_attention', action='store_true', help='incorporate positional information in key/value')
parser.add_argument('--no_source', action='store_true')
parser.add_argument('--use_mask', action='store_true', help='use src/trg mask during attention')
parser.add_argument('--diag', action='store_true', help='ignore diagonal attention when doing self-attention.')
parser.add_argument('--convblock', action='store_true', help='use ConvBlock instead of ResNet')
parser.add_argument('--cosine_output', action='store_true', help='use cosine similarity as output layer')
parser.add_argument('--noisy', action='store_true', help='inject noise in the attention mechanism: Beta-Gumbel softmax')
parser.add_argument('--noise_samples', type=int, default=0, help='only useful for noisy parallel decoding')
parser.add_argument('--critic', action='store_true', help='use critic')
parser.add_argument('--kernel_sizes', type=str, default='2,3,4,5', help='kernel sizes of convnet critic')
parser.add_argument('--kernel_num', type=int, default=128, help='number of each kind of kernel')
parser.add_argument('--use_wo', action='store_true', help='use output weight matrix in multihead attention')
parser.add_argument('--share_embeddings', action='store_true', help='share embeddings between encoder and decoder')
parser.add_argument('--use_alignment', action='store_true', help='use the aligned fake data to initialize')
parser.add_argument('--hard_inputs', action='store_true', help='use hard selection as inputs, instead of soft-attention over embeddings.')
parser.add_argument('--preordering', action='store_true', help='use the ground-truth reordering information')
parser.add_argument('--use_posterior_order', action='store_true', help='directly use the groud-truth alignment for reordering.')
parser.add_argument('--train_decoder_with_order', action='store_true', help='when training the decoder, use the ground-truth')
parser.add_argument('--postordering', action='store_true', help='just have a try...')
parser.add_argument('--fertility_only', action='store_true')
parser.add_argument('--highway', action='store_true', help='usually false')
parser.add_argument('--mix_of_experts', action='store_true')
parser.add_argument('--orderless', action='store_true', help='for the inputs, remove the order information')
parser.add_argument('--cheating', action='store_true', help='disable decoding, always use real fertility')
# running
parser.add_argument('--mode', type=str, default='train', help='train, test or build')
parser.add_argument('--gpu', type=int, default=0, help='GPU to use or -1 for CPU')
parser.add_argument('--seed', type=int, default=19920206, help='seed for randomness')
parser.add_argument('--eval-every', type=int, default=1000, help='run dev every')
parser.add_argument('--maximum_steps', type=int, default=1000000, help='maximum steps you take to train a model')
parser.add_argument('--disable_lr_schedule', action='store_true', help='disable the transformer learning rate')
parser.add_argument('--batchsize', type=int, default=2048, help='# of tokens processed per batch')
parser.add_argument('--hidden_size', type=int, default=None, help='input the hidden size')
parser.add_argument('--length_ratio', type=int, default=2, help='maximum lengths of decoding')
parser.add_argument('--optimizer', type=str, default='Adam')
parser.add_argument('--beam_size', type=int, default=1, help='beam-size used in Beamsearch, default using greedy decoding')
parser.add_argument('--alpha', type=float, default=0.6, help='length normalization weights')
parser.add_argument('--temperature', type=float, default=1, help='smoothing temperature for noisy decoding')
parser.add_argument('--multi_run', type=int, default=1, help='we can run the code multiple times to get the best')
parser.add_argument('--load_from', type=str, default=None, help='load from checkpoint')
parser.add_argument('--resume', action='store_true', help='when loading from the saved model, it resumes from that.')
parser.add_argument('--teacher', type=str, default=None, help='load a pre-trained auto-regressive model.')
parser.add_argument('--share_encoder', action='store_true', help='use teacher-encoder to initialize student')
parser.add_argument('--finetune_encoder', action='store_true', help='if further train the encoder')
parser.add_argument('--seq_dist', action='store_true', help='knowledge distillation at sequence level')
parser.add_argument('--word_dist', action='store_true', help='knowledge distillation at word level')
parser.add_argument('--greedy_fertility', action='store_true', help='using the fertility generated by autoregressive model (only for seq_dist)')
parser.add_argument('--fertility_mode', type=str, default='argmax', help='mean, argmax or reinforce')
parser.add_argument('--finetuning_truth', action='store_true', help='use ground-truth for finetuning')
parser.add_argument('--trainable_teacher', action='store_true', help='have a trainable teacher')
parser.add_argument('--only_update_errors', action='store_true', help='have a trainable teacher')
parser.add_argument('--teacher_use_real', action='store_true', help='teacher also trained with MLE on real data')
parser.add_argument('--max_cache', type=int, default=0, help='save most recent max_cache decoded translations')
parser.add_argument('--replay_every', type=int, default=1000, help='every 1k updates, train the teacher again')
parser.add_argument('--replay_times', type=int, default=250, help='train the teacher again for 250k steps')
parser.add_argument('--margin', type=float, default=1.5, help='margin to make sure teacher will give higher score to real data')
parser.add_argument('--real_data', action='store_true', help='only used in the reverse kl setting')
parser.add_argument('--beta1', type=float, default=0.5, help='balancing MLE and KL loss.')
parser.add_argument('--beta2', type=float, default=0.01, help='balancing the GAN loss.')
parser.add_argument('--critic_only', type=int, default=0, help='pre-training the critic model.')
parser.add_argument('--st', action='store_true', help='straight through estimator')
parser.add_argument('--entropy', action='store_true')
parser.add_argument('--no_bpe', action='store_true', help='output files without BPE')
parser.add_argument('--no_write', action='store_true', help='do not write the decoding into the decoding files.')
parser.add_argument('--output_fer', action='store_true', help='decoding and output fertilities')
# debugging
parser.add_argument('--check', action='store_true', help='on training, only used to check on the test set.')
parser.add_argument('--debug', action='store_true', help='debug mode: no saving or tensorboard')
parser.add_argument('--tensorboard', action='store_true', help='use TensorBoard')
# old params
parser.add_argument('--old', action='store_true', help='this is used for solving conflicts of new codes')
parser.add_argument('--hyperopt', action='store_true', help='use HyperOpt')
parser.add_argument('--scst', action='store_true', help='use HyperOpt')
parser.add_argument('--serve', type=int, default=None, help='serve at port')
parser.add_argument('--attention_discrimination', action='store_true')
# ---------------------------------------------------------------------------------------------------------------- #
args = parser.parse_args()
if args.prefix == '[time]':
args.prefix = strftime("%m.%d_%H.%M.", gmtime())
args.kernel_sizes = [int(k) for k in args.kernel_sizes.split(',')]
# get the langauage pairs:
args.src = args.language[:2] # source language
args.trg = args.language[2:] # target language
# logger settings
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s %(levelname)s: - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
fh = logging.FileHandler('./logs/log-{}.txt'.format(args.prefix))
fh.setLevel(logging.DEBUG)
fh.setFormatter(formatter)
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
ch.setFormatter(formatter)
logger.addHandler(ch)
logger.addHandler(fh)
# setup random seeds
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
# setup data-field
DataField = data.ReversibleField if args.use_revtok else NormalField
tokenizer = revtok.tokenize if args.use_revtok else lambda x: x.replace('@@ ', '').split()
TRG = DataField(init_token='<init>', eos_token='<eos>', batch_first=True)
SRC = DataField(batch_first=True) if not args.share_embeddings else TRG
ALIGN = data.Field(sequential=True, preprocessing=data.Pipeline(lambda tok: int(tok.split('-')[0])), use_vocab=False, pad_token=0, batch_first=True)
FER = data.Field(sequential=True, preprocessing=data.Pipeline(lambda tok: int(tok)), use_vocab=False, pad_token=0, batch_first=True)
align_dict, align_table = None, None
# setup many datasets (need to manaually setup)
data_prefix = args.data_prefix
if args.dataset == 'iwslt':
if args.test_set is None:
args.test_set = 'IWSLT16.TED.tst2013'
if args.dist_set is None:
args.dist_set = '.dec.b1'
elif args.greedy_fertility:
logger.info('use the fertility predicted by autoregressive model (instead of fast-align)')
train_data, dev_data = ParallelDataset.splits(
path=data_prefix + 'iwslt/en-de/', train='train.en-de.bpe.new',
validation='IWSLT16.TED.tst2013.en-de.bpe.new.dev', exts=('.src.b1', '.trg.b1', '.dec.b1', '.fer', '.fer'),
fields=[('src', SRC), ('trg', TRG), ('dec', TRG), ('fer', FER), ('fer_dec', FER)],
load_dataset=args.load_dataset, prefix='ts')
elif (args.mode == 'test') or (args.mode == 'test_noisy'):
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'iwslt/en-de/', train='train.tags.en-de{}'.format(
'.bpe' if not args.use_revtok else ''),
validation='{}.en-de{}'.format(
args.test_set, '.bpe' if not args.use_revtok else ''), exts=('.en', '.de'),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='normal')
else:
train_data, dev_data = ParallelDataset.splits(
path=data_prefix + 'iwslt/en-de/', train='train.tags.en-de.bpe',
validation='train.tags.en-de.bpe.dev', exts=('.en2', '.de2', '.decoded2', '.aligned', '.decode.aligned', '.fer', '.decode.fer'),
fields=[('src', SRC), ('trg', TRG), ('dec', TRG), ('align', ALIGN), ('align_dec', ALIGN), ('fer', FER), ('fer_dec', FER)],
load_dataset=args.load_dataset, prefix='ts')
decoding_path = data_prefix + 'iwslt/en-de/{}.en-de.bpe.new'
if args.use_alignment and (args.model is FastTransformer):
align_dict = {l.split()[0]: l.split()[1] for l in open(data_prefix + 'iwslt/en-de/train.tags.en-de.dict')}
elif args.dataset == 'wmt16-ende':
if args.test_set is None:
args.test_set = 'newstest2013'
if (args.mode == 'test') or (args.mode == 'test_noisy'):
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'wmt16-ende/', train='newstest2013.tok.bpe.32000',
validation='{}.tok.bpe.32000'.format(args.test_set), exts=('.{}'.format(args.src), '.{}'.format(args.trg)),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='real')
decoding_path = data_prefix + 'wmt16-ende/test.{}.{}'.format(args.prefix, args.test_set)
elif not args.seq_dist:
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'wmt16-ende/', train='train.tok.clean.bpe.32000',
validation='{}.tok.bpe.32000'.format(args.test_set), exts=('.{}'.format(args.src), '.{}'.format(args.trg)),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='real')
decoding_path = data_prefix + 'wmt16-ende/{}.tok.bpe.decode'
else:
train_data, dev_data = ParallelDataset.splits(
path=data_prefix + 'wmt16-ende/', train='train.tok.bpe.decode',
validation='newstest2013.tok.bpe.decode.dev',
exts=('.src.b1', '.trg.b1', '.dec.b1', '.real.aligned', '.fake.aligned', '.real.fer', '.fake.fer'),
fields=[('src', SRC), ('trg', TRG), ('dec', TRG), ('align', ALIGN), ('align_dec', ALIGN), ('fer', FER), ('fer_dec', FER)],
load_dataset=args.load_dataset, prefix='ts')
decoding_path = data_prefix + 'wmt16-ende/{}.tok.bpe.na'
if args.use_alignment and (args.model is FastTransformer):
align_table = {l.split()[0]: l.split()[1] for l in
open(data_prefix + 'wmt16-ende/train.tok.bpe.decode.full.fastlign2.dict')}
elif args.dataset == 'wmt16-deen':
if args.test_set is None:
args.test_set = 'newstest2013'
if (args.mode == 'test') or (args.mode == 'test_noisy'):
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'wmt16-ende/', train='newstest2013.tok.bpe.32000',
validation='{}.tok.bpe.32000'.format(args.test_set), exts=('.{}'.format(args.src), '.{}'.format(args.trg)),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='real')
decoding_path = data_prefix + 'wmt16-ende/test.{}.{}'.format(args.prefix, args.test_set)
elif not args.seq_dist:
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'wmt16-deen/', train='train.tok.clean.bpe.32000',
validation='{}.tok.bpe.32000'.format(args.test_set), exts=('.{}'.format(args.src), '.{}'.format(args.trg)),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='real')
decoding_path = data_prefix + 'wmt16-deen/{}.tok.bpe.decode'
else:
train_data, dev_data = ParallelDataset.splits(
path=data_prefix + 'wmt16-deen/', train='train.tok.bpe.decode',
validation='{}.tok.bpe.decode.dev'.format(args.test_set),
exts=('.src.b1', '.trg.b1', '.dec.b1', '.real.aligned', '.fake.aligned', '.real.fer', '.fake.fer'),
fields=[('src', SRC), ('trg', TRG), ('dec', TRG), ('align', ALIGN), ('align_dec', ALIGN), ('fer', FER), ('fer_dec', FER)],
load_dataset=args.load_dataset, prefix='ts')
decoding_path = data_prefix + 'wmt16-deen/{}.tok.bpe.na'
if args.use_alignment and (args.model is FastTransformer):
align_table = {l.split()[0]: l.split()[1] for l in
open(data_prefix + 'wmt16-deen/train.tok.bpe.decode.full.fastlign2.dict')}
elif args.dataset == 'wmt16-enro':
if args.test_set is None:
args.test_set = 'dev'
if (args.mode == 'test') or (args.mode == 'test_noisy'):
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'wmt16-enro/', train='dev.bpe',
validation='{}.bpe'.format(args.test_set), exts=('.{}'.format(args.src), '.{}'.format(args.trg)),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='real')
decoding_path = data_prefix + 'wmt16-enro/{}.bpe.decode'
elif not args.seq_dist:
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'wmt16-enro/', train='corpus.bpe',
validation='{}.bpe'.format(args.test_set), exts=('.{}'.format(args.src), '.{}'.format(args.trg)),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='real')
decoding_path = data_prefix + 'wmt16-enro/{}.bpe.decode'
else:
train_data, dev_data = ParallelDataset.splits(
path=data_prefix + 'wmt16-enro/', train='train.bpe.decode',
validation='dev.bpe.decode.dev',
exts=('.src.b1', '.trg.b1', '.dec.b1', '.real.aligned', '.fake.aligned', '.real.fer', '.fake.fer'),
fields=[('src', SRC), ('trg', TRG), ('dec', TRG), ('align', ALIGN), ('align_dec', ALIGN), ('fer', FER), ('fer_dec', FER)],
load_dataset=args.load_dataset, prefix='ts')
decoding_path = data_prefix + 'wmt16-enro/{}.tok.bpe.na'
if args.use_alignment and (args.model is FastTransformer):
align_table = {l.split()[0]: l.split()[1] for l in
open(data_prefix + 'wmt16-enro/train.bpe.decode.full.fastlign2.dict')}
elif args.dataset == 'wmt16-roen':
if args.test_set is None:
args.test_set = 'dev'
if (args.mode == 'test') or (args.mode == 'test_noisy'):
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'wmt16-roen/', train='dev.bpe',
validation='{}.bpe'.format(args.test_set), exts=('.{}'.format(args.src), '.{}'.format(args.trg)),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='real')
decoding_path = data_prefix + 'wmt16-roen/{}.bpe.decode'
elif not args.seq_dist:
train_data, dev_data = NormalTranslationDataset.splits(
path=data_prefix + 'wmt16-roen/', train='corpus.bpe',
validation='{}.bpe'.format(args.test_set), exts=('.{}'.format(args.src), '.{}'.format(args.trg)),
fields=(SRC, TRG), load_dataset=args.load_dataset, prefix='real')
decoding_path = data_prefix + 'wmt16-roen/{}.bpe.decode'
else:
train_data, dev_data = ParallelDataset.splits(
path=data_prefix + 'wmt16-roen/', train='train.bpe.decode',
validation='dev.bpe.decode.dev',
exts=('.src.b1', '.trg.b1', '.dec.b1', '.real.aligned', '.fake.aligned', '.real.fer', '.fake.fer'),
fields=[('src', SRC), ('trg', TRG), ('dec', TRG), ('align', ALIGN), ('align_dec', ALIGN), ('fer', FER), ('fer_dec', FER)],
load_dataset=args.load_dataset, prefix='ts')
decoding_path = data_prefix + 'wmt16-roen/{}.tok.bpe.na'
if args.use_alignment and (args.model is FastTransformer):
align_table = {l.split()[0]: l.split()[1] for l in
open(data_prefix + 'wmt16-roen/train.bpe.decode.full.fastlign2.dict')}
else:
raise NotImplementedError
# build word-level vocabularies
if args.load_vocab and os.path.exists(data_prefix + '{}/vocab{}_{}.pt'.format(
args.dataset, 'shared' if args.share_embeddings else '', '{}-{}'.format(args.src, args.trg))):
logger.info('load saved vocabulary.')
src_vocab, trg_vocab = torch.load(data_prefix + '{}/vocab{}_{}.pt'.format(
args.dataset, 'shared' if args.share_embeddings else '', '{}-{}'.format(args.src, args.trg)))
SRC.vocab = src_vocab
TRG.vocab = trg_vocab
else:
logger.info('save the vocabulary')
if not args.share_embeddings:
SRC.build_vocab(train_data, dev_data, max_size=50000)
TRG.build_vocab(train_data, dev_data, max_size=50000)
torch.save([SRC.vocab, TRG.vocab], data_prefix + '{}/vocab{}_{}.pt'.format(
args.dataset, 'shared' if args.share_embeddings else '', '{}-{}'.format(args.src, args.trg)))
args.__dict__.update({'trg_vocab': len(TRG.vocab), 'src_vocab': len(SRC.vocab)})
# build alignments ---
if align_dict is not None:
align_table = [TRG.vocab.stoi['<init>'] for _ in range(len(SRC.vocab.itos))]
for src in align_dict:
align_table[SRC.vocab.stoi[src]] = TRG.vocab.stoi[align_dict[src]]
align_table[0] = 0 # --<unk>
align_table[1] = 1 # --<pad>
def dyn_batch_with_padding(new, i, sofar):
prev_max_len = sofar / (i - 1) if i > 1 else 0
if args.seq_dist:
return max(len(new.src), len(new.trg), len(new.dec), prev_max_len) * i
else:
return max(len(new.src), len(new.trg), prev_max_len) * i
def dyn_batch_without_padding(new, i, sofar):
if args.seq_dist:
return sofar + max(len(new.src), len(new.trg), len(new.dec))
else:
return sofar + max(len(new.src), len(new.trg))
# build the dataset iterators
# work around torchtext making it hard to share vocabs without sharing other field properties
if args.share_embeddings:
SRC = copy.deepcopy(SRC)
SRC.init_token = None
SRC.eos_token = None
train_data.fields['src'] = SRC
dev_data.fields['src'] = SRC
if (args.model is FastTransformer) and (args.remove_eos):
TRG.eos_token = None
if args.max_len is not None:
train_data.examples = [ex for ex in train_data.examples if len(ex.trg) <= args.max_len]
if args.batchsize == 1: # speed-test: one sentence per batch.
batch_size_fn = lambda new, count, sofar: count
else:
batch_size_fn = dyn_batch_without_padding if args.model is Transformer else dyn_batch_with_padding
train_real, dev_real = data.BucketIterator.splits(
(train_data, dev_data), batch_sizes=(args.batchsize, args.batchsize), device=args.gpu,
batch_size_fn=batch_size_fn,
repeat=None if args.mode == 'train' else False)
logger.info("build the dataset. done!")
# model hyper-params:
hparams = None
if args.dataset == 'iwslt':
if args.params == 'james-iwslt':
hparams = {'d_model': 278, 'd_hidden': 507, 'n_layers': 5,
'n_heads': 2, 'drop_ratio': 0.079, 'warmup': 746} # ~32
elif args.params == 'james-iwslt2':
hparams = {'d_model': 278, 'd_hidden': 2048, 'n_layers': 5,
'n_heads': 2, 'drop_ratio': 0.079, 'warmup': 746} # ~32
teacher_hparams = {'d_model': 278, 'd_hidden': 507, 'n_layers': 5,
'n_heads': 2, 'drop_ratio': 0.079, 'warmup': 746}
elif args.dataset == 'wmt16-ende':
logger.info('use default parameters of t2t-base')
hparams = {'d_model': 512, 'd_hidden': 512, 'n_layers': 6,
'n_heads': 8, 'drop_ratio': 0.1, 'warmup': 16000} # ~32
teacher_hparams = hparams
elif args.dataset == 'wmt16-deen':
logger.info('use default parameters of t2t-base')
hparams = {'d_model': 512, 'd_hidden': 512, 'n_layers': 6,
'n_heads': 8, 'drop_ratio': 0.1, 'warmup': 16000} # ~32
teacher_hparams = hparams
elif args.dataset == 'wmt16-enro':
logger.info('use default parameters of t2t-base')
hparams = {'d_model': 512, 'd_hidden': 512, 'n_layers': 6,
'n_heads': 8, 'drop_ratio': 0.1, 'warmup': 16000} # ~32
teacher_hparams = hparams
elif args.dataset == 'wmt16-roen':
logger.info('use default parameters of t2t-base')
hparams = {'d_model': 512, 'd_hidden': 512, 'n_layers': 6,
'n_heads': 8, 'drop_ratio': 0.1, 'warmup': 16000} # ~32
teacher_hparams = hparams
if hparams is None:
logger.info('use default parameters of t2t-base')
hparams = {'d_model': 512, 'd_hidden': 512, 'n_layers': 6,
'n_heads': 8, 'drop_ratio': 0.1, 'warmup': 16000} # ~32
if args.teacher is not None:
teacher_args = copy.deepcopy(args)
teacher_args.__dict__.update(teacher_hparams)
args.__dict__.update(hparams)
if args.hidden_size is not None:
args.d_hidden = args.hidden_size
# show the arg:
logger.info(args)
hp_str = (f"{args.dataset}_{args.level}_{'fast_' if args.model is FastTransformer else ''}"
f"{args.d_model}_{args.d_hidden}_{args.n_layers}_{args.n_heads}_"
f"{args.drop_ratio:.3f}_{args.warmup}_"
f"{args.xe_until if hasattr(args, 'xe_until') else ''}_"
f"{f'{args.xe_ratio:.3f}' if hasattr(args, 'xe_ratio') else ''}_"
f"{args.xe_every if hasattr(args, 'xe_every') else ''}")
logger.info(f'Starting with HPARAMS: {hp_str}')
model_name = './models/' + args.prefix + hp_str
# build the model
model = args.model(SRC, TRG, args)
if args.load_from is not None:
with torch.cuda.device(args.gpu): # very important.
model.load_state_dict(torch.load('./models/' + args.load_from + '.pt',
map_location=lambda storage, loc: storage.cuda())) # load the pretrained models.
if args.critic:
model.install_critic()
# logger.info(str(model))
# if using a teacher
if args.teacher is not None:
teacher_model = Transformer(SRC, TRG, teacher_args)
with torch.cuda.device(args.gpu):
teacher_model.load_state_dict(torch.load('./models/' + args.teacher + '.pt',
map_location=lambda storage, loc: storage.cuda()))
for params in teacher_model.parameters():
if args.trainable_teacher:
params.requires_grad = True
else:
params.requires_grad = False
if (args.share_encoder) and (args.load_from is None):
model.encoder = copy.deepcopy(teacher_model.encoder)
for params in model.encoder.parameters():
if args.finetune_encoder:
params.requires_grad = True
else:
params.requires_grad = False
else:
teacher_model = None
# use cuda
if args.gpu > -1:
model.cuda(args.gpu)
if align_table is not None:
align_table = torch.LongTensor(align_table).cuda(args.gpu)
align_table = Variable(align_table)
model.alignment = align_table
if args.teacher is not None:
teacher_model.cuda(args.gpu)
def register_nan_checks(m):
def check_grad(module, grad_input, grad_output):
if any(np.any(np.isnan(gi.data.cpu().numpy())) for gi in grad_input if gi is not None):
print('NaN gradient in ' + type(module).__name__)
1/0
m.apply(lambda module: module.register_backward_hook(check_grad))
def get_learning_rate(i, lr0=0.1):
if not args.disable_lr_schedule:
return lr0 * 10 / math.sqrt(args.d_model) * min(
1 / math.sqrt(i), i / (args.warmup * math.sqrt(args.warmup)))
return 0.00002
def export(x):
try:
with torch.cuda.device(args.gpu):
return x.data.cpu().float().mean()
except Exception:
return 0
def devol(batch):
new_batch = copy.copy(batch)
new_batch.src = Variable(batch.src.data, volatile=True)
return new_batch
# register_nan_checks(model)
# register_nan_checks(teacher_model)
def valid_model(model, dev, dev_metrics=None, distillation=False, print_out=False, teacher_model=None):
print_seqs = ['[sources]', '[targets]', '[decoded]', '[fertili]', '[origind]']
trg_outputs, dec_outputs = [], []
outputs = {}
model.eval()
if teacher_model is not None:
teacher_model.eval()
for j, dev_batch in enumerate(dev):
# decode from the model (whatever Transformer or FastTransformer)
torch.cuda.nvtx.range_push('quick_prepare')
inputs, input_masks, targets, target_masks, sources, source_masks, encoding, batch_size = model.quick_prepare(dev_batch, distillation)
torch.cuda.nvtx.range_pop()
torch.cuda.nvtx.range_push('prepare_initial')
decoder_inputs, input_reorder, reordering_cost = inputs, None, None
if type(model) is FastTransformer:
# batch_align = dev_batch.align_dec if distillation else dev_batch.align
batch_align = None
batch_fer = dev_batch.fer_dec if distillation else dev_batch.fer
# if args.postordering:
#
# targets_sorted = targets.gather(1, align_index)
# batch_align_sorted, align_index = masked_sort(batch_align, target_masks) # change the target indexxx, batch x max_trg
decoder_inputs, input_reorder, decoder_masks, reordering_cost = model.prepare_initial(encoding,
sources, source_masks, input_masks,
batch_align, batch_fer, decoding=(not args.cheating), mode='argmax')
else:
decoder_masks = input_masks
torch.cuda.nvtx.range_pop()
torch.cuda.nvtx.range_push('model')
decoding, out, probs = model(encoding, source_masks, decoder_inputs, decoder_masks, decoding=True, return_probs=True)
torch.cuda.nvtx.range_pop()
torch.cuda.nvtx.range_push('batched_cost')
loss = 0
if args.postordering:
if args.cheating:
decoding1 = unsorted(decoding, align_index)
else:
positions = model.predict_offset(out, decoder_masks, None)
shifted_index = positions.sort(1)[1]
decoding1 = unsorted(decoding, shifted_index)
else:
decoding1 = decoding
# loss = model.batched_cost(targets, target_masks, probs)
torch.cuda.nvtx.range_pop()
torch.cuda.nvtx.range_push('output_decoding')
dev_outputs = [model.output_decoding(d) for d in [('src', sources), ('trg', targets), ('trg', decoding1), ('src', input_reorder)]]
if args.postordering:
dev_outputs += [model.output_decoding(('trg', decoding))]
torch.cuda.nvtx.range_pop()
torch.cuda.nvtx.range_push('computeGLEU')
gleu = computeGLEU(dev_outputs[2], dev_outputs[1], corpus=False, tokenizer=tokenizer)
torch.cuda.nvtx.range_pop()
if print_out:
for k, d in enumerate(dev_outputs):
logger.info("{}: {}".format(print_seqs[k], d[0]))
logger.info('------------------------------------------------------------------')
if teacher_model is not None: # teacher is Transformer, student is FastTransformer
inputs_student, _, targets_student, _, _, _, encoding_teacher, _ = teacher_model.quick_prepare(dev_batch, False, decoding, decoding,
input_masks, target_masks, source_masks)
teacher_real_loss = teacher_model.cost(targets, target_masks,
out=teacher_model(encoding_teacher, source_masks, inputs, input_masks))
teacher_fake_out = teacher_model(encoding_teacher, source_masks, inputs_student, input_masks)
teacher_fake_loss = teacher_model.cost(targets_student, target_masks, out=teacher_fake_out)
teacher_alter_loss = teacher_model.cost(targets, target_masks, out=teacher_fake_out)
trg_outputs += dev_outputs[1]
dec_outputs += dev_outputs[2]
if dev_metrics is not None:
values = [loss, gleu]
if teacher_model is not None:
values += [teacher_real_loss, teacher_fake_loss,
teacher_real_loss - teacher_fake_loss,
teacher_alter_loss,
teacher_alter_loss - teacher_fake_loss]
if reordering_cost is not None:
values += [reordering_cost]
dev_metrics.accumulate(batch_size, *values)
corpus_gleu = computeGLEU(dec_outputs, trg_outputs, corpus=True, tokenizer=tokenizer)
corpus_bleu = computeBLEU(dec_outputs, trg_outputs, corpus=True, tokenizer=tokenizer)
outputs['corpus_gleu'] = corpus_gleu
outputs['corpus_bleu'] = corpus_bleu
if dev_metrics is not None:
logger.info(dev_metrics)
logger.info("The dev-set corpus GLEU = {}".format(corpus_gleu))
logger.info("The dev-set corpus BLEU = {}".format(corpus_bleu))
return outputs
def train_model(model, train, dev, teacher_model=None):
if args.tensorboard and (not args.debug):
from tensorboardX import SummaryWriter
writer = SummaryWriter('./runs/{}'.format(args.prefix+hp_str))
# optimizer
if args.optimizer == 'Adam':
opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], betas=(0.9, 0.98), eps=1e-9)
if args.trainable_teacher:
opt_teacher = torch.optim.Adam([p for p in teacher_model.parameters() if p.requires_grad], betas=(0.9, 0.98), eps=1e-9)
elif args.optimizer == 'RMSprop':
opt = torch.optim.RMSprop([p for p in model.parameters() if p.requires_grad], eps=1e-9)
if args.trainable_teacher:
opt_teacher = torch.optim.RMSprop([p for p in teacher_model.parameters() if p.requires_grad], eps=1e-9)
else:
raise NotImplementedError
# if resume training
if (args.load_from is not None) and (args.resume):
with torch.cuda.device(args.gpu): # very important.
offset, opt_states = torch.load('./models/' + args.load_from + '.pt.states',
map_location=lambda storage, loc: storage.cuda())
opt.load_state_dict(opt_states)
else:
offset = 0
# metrics
best = Best(max, 'corpus_bleu', 'corpus_gleu', 'gleu', 'loss', 'i', model=model, opt=opt, path=model_name, gpu=args.gpu)
train_metrics = Metrics('train', 'loss', 'real', 'fake')
dev_metrics = Metrics('dev', 'loss', 'gleu', 'real_loss', 'fake_loss', 'distance', 'alter_loss', 'distance2', 'reordering_loss', 'corpus_gleu')
progressbar = tqdm(total=args.eval_every, desc='start training.')
# cache
if args.max_cache > 0:
caches = Cache(args.max_cache, args.gpu)
for iters, batch in enumerate(train):
iters += offset
if iters > args.maximum_steps:
logger.info('reach the maximum updating steps.')
break
if iters % args.eval_every == 0:
progressbar.close()
dev_metrics.reset()
if args.seq_dist:
outputs_course = valid_model(model, dev, dev_metrics,
distillation=True, teacher_model=None)#teacher_model=teacher_model)
if args.trainable_teacher:
outputs_teacher = valid_model(teacher_model, dev, None)
outputs_data = valid_model(model, dev, None if args.seq_dist else dev_metrics, teacher_model=None, print_out=True)
if args.tensorboard and (not args.debug):
writer.add_scalar('dev/GLEU_sentence_', dev_metrics.gleu, iters)
writer.add_scalar('dev/Loss', dev_metrics.loss, iters)
writer.add_scalar('dev/GLEU_corpus_', outputs_data['corpus_gleu'], iters)
writer.add_scalar('dev/BLEU_corpus_', outputs_data['corpus_bleu'], iters)
if args.seq_dist:
writer.add_scalar('dev/GLEU_corpus_dis', outputs_course['corpus_gleu'], iters)
writer.add_scalar('dev/BLEU_corpus_dis', outputs_course['corpus_bleu'], iters)
if args.trainable_teacher:
writer.add_scalar('dev/GLEU_corpus_teacher', outputs_teacher['corpus_gleu'], iters)
writer.add_scalar('dev/BLEU_corpus_teacher', outputs_teacher['corpus_bleu'], iters)
if args.teacher is not None:
writer.add_scalar('dev/Teacher_real_loss', dev_metrics.real_loss, iters)
writer.add_scalar('dev/Teacher_fake_loss', dev_metrics.fake_loss, iters)
writer.add_scalar('dev/Teacher_alter_loss', dev_metrics.alter_loss, iters)
writer.add_scalar('dev/Teacher_distance', dev_metrics.distance, iters)
writer.add_scalar('dev/Teacher_distance2', dev_metrics.distance2, iters)
if args.preordering:
writer.add_scalar('dev/Reordering_loss', dev_metrics.reordering_loss, iters)
if not args.debug:
best.accumulate(outputs_data['corpus_bleu'], outputs_data['corpus_gleu'], dev_metrics.gleu, dev_metrics.loss, iters)
logger.info('the best model is achieved at {}, average greedy GLEU={}, corpus GLEU={}, corpus BLEU={}'.format(
best.i, best.gleu, best.corpus_gleu, best.corpus_bleu))
logger.info('model:' + args.prefix + hp_str)
# ---set-up a new progressor---
progressbar = tqdm(total=args.eval_every, desc='start training.')
# --- training --- #
# try:
model.train()
opt.param_groups[0]['lr'] = get_learning_rate(iters + 1)
opt.zero_grad()
# prepare the data
inputs, input_masks, targets, target_masks, sources, source_masks, encoding, batch_size = model.quick_prepare(batch, args.seq_dist)
input_reorder, reordering_cost, decoder_inputs = None, None, inputs
batch_align = None # batch.align_dec if args.seq_dist else batch.align
batch_fer = batch.fer_dec if args.seq_dist else batch.fer
# batch_align_sorted, align_index = masked_sort(batch_align, target_masks) # change the target indexxx, batch x max_trg
# print(batch_fer.size(), input_masks.size(), source_masks.size(), sources.size())
# Prepare_Initial
if type(model) is FastTransformer:
inputs, input_reorder, input_masks, reordering_cost = model.prepare_initial(encoding, sources, source_masks, input_masks, batch_align, batch_fer)
# Maximum Likelihood Training
feedback = {}
if not args.word_dist:
loss = model.cost(targets, target_masks, out=model(encoding, source_masks, inputs, input_masks, positions= None, feedback=feedback))
# train the reordering also using MLE??
if args.preordering:
loss += reordering_cost
else:
# only used for FastTransformer: word-level adjustment
if not args.preordering:
decoding, out, probs = model(encoding, source_masks, inputs, input_masks, return_probs=True, decoding=True)
loss_student = model.batched_cost(targets, target_masks, probs) # student-loss (MLE)
decoder_masks = input_masks
else: # Note that MLE and decoding has different translations. We need to run the same code twice
if args.finetuning_truth:
decoding, out, probs = model(encoding, source_masks, inputs, input_masks, decoding=True, return_probs=True, feedback=feedback)
loss_student = model.cost(targets, target_masks, out=out)
decoder_masks = input_masks
else:
if args.fertility_mode != 'reinforce':
loss_student = model.cost(targets, target_masks, out=model(encoding, source_masks, inputs, input_masks, positions=None, feedback=feedback))
decoder_inputs, _, decoder_masks, _ = model.prepare_initial(encoding, sources, source_masks, input_masks,
batch_align, batch_fer, decoding=True, mode=args.fertility_mode)
decoding, out, probs = model(encoding, source_masks, decoder_inputs, decoder_masks, decoding=True, return_probs=True) # decode again
else:
# truth
decoding, out, probs = model(encoding, source_masks, inputs, input_masks, decoding=True, return_probs=True, feedback=feedback)
loss_student = model.cost(targets, target_masks, out=out)
decoder_masks = input_masks
# baseline
decoder_inputs_b, _, decoder_masks_b, _ = model.prepare_initial(encoding, sources, source_masks, input_masks,
batch_align, batch_fer, decoding=True, mode='mean')
decoding_b, out_b, probs_b = model(encoding, source_masks, decoder_inputs_b, decoder_masks_b, decoding=True, return_probs=True) # decode again
# reinforce
decoder_inputs_r, _, decoder_masks_r, _ = model.prepare_initial(encoding, sources, source_masks, input_masks,
batch_align, batch_fer, decoding=True, mode='reinforce')
decoding_r, out_r, probs_r = model(encoding, source_masks, decoder_inputs_r, decoder_masks_r, decoding=True, return_probs=True) # decode again
# train the reordering also using MLE??
if args.preordering:
loss_student += reordering_cost
# teacher tries translation + look-at student's output
teacher_model.eval()
if args.fertility_mode != 'reinforce':
inputs_student_index, _, targets_student_soft, _, _, _, encoding_teacher, _ = model.quick_prepare(batch, False, decoding, probs, decoder_masks, decoder_masks, source_masks)
out_teacher, probs_teacher = teacher_model(encoding_teacher, source_masks, inputs_student_index.detach(), decoder_masks, return_probs=True)
loss_teacher = teacher_model.batched_cost(targets_student_soft, decoder_masks, probs_teacher.detach())
loss = (1 - args.beta1) * loss_teacher + args.beta1 * loss_student # final results
else:
inputs_student_index, _, targets_student_soft, _, _, _, encoding_teacher, _ = model.quick_prepare(batch, False, decoding, probs, decoder_masks, decoder_masks, source_masks)
out_teacher, probs_teacher = teacher_model(encoding_teacher, source_masks, inputs_student_index.detach(), decoder_masks, return_probs=True)
loss_teacher = teacher_model.batched_cost(targets_student_soft, decoder_masks, probs_teacher.detach())
inputs_student_index, _ = model.prepare_inputs(batch, decoding_b, False, decoder_masks_b)
targets_student_soft, _ = model.prepare_targets(batch, probs_b, False, decoder_masks_b)
out_teacher, probs_teacher = teacher_model(encoding_teacher, source_masks, inputs_student_index.detach(), decoder_masks_b, return_probs=True)
_, loss_1= teacher_model.batched_cost(targets_student_soft, decoder_masks_b, probs_teacher.detach(), True)
inputs_student_index, _ = model.prepare_inputs(batch, decoding_r, False, decoder_masks_r)
targets_student_soft, _ = model.prepare_targets(batch, probs_r, False, decoder_masks_r)
out_teacher, probs_teacher = teacher_model(encoding_teacher, source_masks, inputs_student_index.detach(), decoder_masks_r, return_probs=True)
_, loss_2= teacher_model.batched_cost(targets_student_soft, decoder_masks_r, probs_teacher.detach(), True)
rewards = -(loss_2 - loss_1).data
# if rewards.size(0) != 1:
rewards = rewards - rewards.mean() # ) / (rewards.std() + TINY)
rewards = rewards.expand_as(source_masks)
rewards = rewards * source_masks
# print(model.predictor.saved_fertilities)
# print(batch.src.size())
model.predictor.saved_fertilities.reinforce(0.1 * rewards.contiguous().view(-1, 1))
loss = (1 - args.beta1) * loss_teacher + args.beta1 * loss_student #+ 0 * model.predictor.saved_fertilities.float().sum() # detect reinforce
# loss = 0 * model.predictor.saved_fertilities.float().sum() # detect reinforce
# accmulate the training metrics
train_metrics.accumulate(batch_size, loss, print_iter=None)
train_metrics.reset()
# train the student
if args.preordering and args.fertility_mode == 'reinforce':
torch.autograd.backward((loss, model.predictor.saved_fertilities),
(torch.ones(1).cuda(loss.get_device()), None))
else:
loss.backward()
# torch.nn.utils.clip_grad_norm(model.parameters(), 1)
opt.step()
info = 'training step={}, loss={:.3f}, lr={:.5f}'.format(iters, export(loss), opt.param_groups[0]['lr'])
if args.word_dist:
info += '| NA:{:.3f}, AR:{:.3f}'.format(export(loss_student), export(loss_teacher))
if args.trainable_teacher and (args.max_cache <= 0):
loss_alter, loss_worse = export(loss_alter), export(loss_worse)
info += '| AL:{:.3f}, WO:{:.3f}'.format(loss_alter, loss_worse)
if args.preordering:
info += '| RE:{:.3f}'.format(export(reordering_cost))
if args.fertility_mode == 'reinforce':
info += '| RL: {:.3f}'.format(export(rewards.mean()))
if args.max_cache > 0:
info += '| caches={}'.format(len(caches.cache))
if args.tensorboard and (not args.debug):
writer.add_scalar('train/Loss', export(loss), iters)
progressbar.update(1)
progressbar.set_description(info)
# continue-training the teacher model
if args.trainable_teacher:
if args.max_cache > 0:
caches.add([batch.src, batch.trg, batch.dec, decoding]) # experience-reply
# trainable teacher: used old experience to train
if (iters+1) % args.replay_every == 0:
# ---set-up a new progressor: teacher training--- #
progressbar_teacher = tqdm(total=args.replay_times, desc='start training the teacher.')
for j in range(args.replay_times):
opt_teacher.param_groups[0]['lr'] = get_learning_rate(iters + 1)
opt_teacher.zero_grad()
src, trg, dec, decoding = caches.sample()
batch = Batch(src, trg, dec)
inputs, input_masks, targets, target_masks, sources, source_masks, encoding_teacher, batch_size = teacher_model.quick_prepare(batch, (not args.teacher_use_real))
inputs_students, _ = teacher_model.prepare_inputs(batch, decoding, masks=input_masks)
loss_alter = teacher_model.cost(targets, target_masks, out=teacher_model(encoding_teacher, source_masks, inputs_students, input_masks))
loss_worse = teacher_model.cost(targets, target_masks, out=teacher_model(encoding_teacher, source_masks, inputs, input_masks))
loss2 = loss_alter + loss_worse
loss2.backward()
opt_teacher.step()
info = 'teacher step={}, loss={:.3f}, alter={:.3f}, worse={:.3f}'.format(j, export(loss2), export(loss_alter), export(loss_worse))
progressbar_teacher.update(1)
progressbar_teacher.set_description(info)
progressbar_teacher.close()
# except Exception as e:
# logger.warn('caught an exception: {}'.format(e))
def decode_model(model, train_real, dev_real, evaluate=True, decoding_path=None, names=['en', 'de', 'decode']):
if train_real is None:
logger.info('decoding from the devlopment set. beamsize={}, alpha={}'.format(args.beam_size, args.alpha))
dev = dev_real
else:
logger.info('decoding from the training set. beamsize={}, alpha={}'.format(args.beam_size, args.alpha))
dev = train_real
dev.train = False # make the Iterator create Variables with volatile=True so no graph is built
progressbar = tqdm(total=sum([1 for _ in dev]), desc='start decoding')
model.eval()
if decoding_path is not None:
decoding_path = decoding_path.format(args.test_set if train_real is None else 'train')
handle_dec = open(decoding_path + '.{}'.format(names[2]), 'w')
handle_src = open(decoding_path + '.{}'.format(names[0]), 'w')
handle_trg = open(decoding_path + '.{}'.format(names[1]), 'w')
if args.output_fer:
handle_fer = open(decoding_path + '.{}'.format('fer'), 'w')
corpus_size = 0
src_outputs, trg_outputs, dec_outputs, timings = [], [], [], []
decoded_words, target_words, decoded_info = 0, 0, 0
attentions = None #{'source': None, 'target': None}
pad_id = model.decoder.field.vocab.stoi['<pad>']
eos_id = model.decoder.field.vocab.stoi['<eos>']
curr_time = 0
for iters, dev_batch in enumerate(dev):
start_t = time.time()
inputs, input_masks, targets, target_masks, sources, source_masks, encoding, batch_size = model.quick_prepare(dev_batch)
if args.model is FastTransformer:
decoder_inputs, input_reorder, decoder_masks, _ = model.prepare_initial(encoding, sources, source_masks, input_masks,
None, None, decoding=True, mode=args.fertility_mode)
else:
decoder_inputs, decoder_masks = inputs, input_masks
decoding = model(encoding, source_masks, decoder_inputs, decoder_masks, beam=args.beam_size, alpha=args.alpha, decoding=True, feedback=attentions)
used_t = time.time() - start_t
curr_time += used_t
real_mask = 1 - ((decoding.data == eos_id) + (decoding.data == pad_id)).float()
outputs = [model.output_decoding(d) for d in [('src', sources), ('trg', targets), ('trg', decoding)]]
def DHondt(approx, mask):
L = mask.size(1)
w = torch.arange(1, 2 * L, 2)
if approx.is_cuda:
w = w.cuda(approx.get_device())
w = 1 / w # 1, 1/2, 1/3, ...
approx = approx[:, :, None] @ w[None, :] # B x Ts x Tt
approx = approx.view(approx.size(0), -1) # B x (Ts x Tt)
appinx = approx.topk(L, 1)[1] # B x Tt (index)
fertility = approx.new(*approx.size()).fill_(0).scatter_(1, appinx, mask)
fertility = fertility.contiguous().view(mask.size(0), -1, mask.size(1)).sum(2).long()
return fertility
def cutoff(s, t):
for i in range(len(s), 0, -1):
if s[i-1] != t:
return s[:i]
raise IndexError
if args.output_fer:
source_attention = attentions['source'].data.mean(1).transpose(2, 1) # B x Ts x Tt
source_attention *= real_mask[:, None, :]
approx_fertility = source_attention.sum(2) # B x Ts
fertility = DHondt(approx_fertility, real_mask)
corpus_size += batch_size
src_outputs += outputs[0]
trg_outputs += outputs[1]
dec_outputs += outputs[2]
timings += [used_t]
if decoding_path is not None:
for s, t, d in zip(outputs[0], outputs[1], outputs[2]):
if args.no_bpe:
s, t, d = s.replace('@@ ', ''), t.replace('@@ ', ''), d.replace('@@ ', '')
print(s, file=handle_src, flush=True)
print(t, file=handle_trg, flush=True)
print(d, file=handle_dec, flush=True)
if args.output_fer:
with torch.cuda.device_of(fertility):
fertility = fertility.tolist()
for f in fertility:
f = ' '.join([str(fi) for fi in cutoff(f, 0)])
print(f, file=handle_fer, flush=True)
progressbar.update(1)
progressbar.set_description('finishing sentences={}/batches={}, speed={} sec/batch'.format(corpus_size, iters, curr_time / (1 + iters)))
if evaluate:
corpus_gleu = computeGLEU(dec_outputs, trg_outputs, corpus=True, tokenizer=tokenizer)
corpus_bleu = computeBLEU(dec_outputs, trg_outputs, corpus=True, tokenizer=tokenizer)
logger.info("The dev-set corpus GLEU = {}".format(corpus_gleu))
logger.info("The dev-set corpus BLEU = {}".format(corpus_bleu))
computeGroupBLEU(dec_outputs, trg_outputs, tokenizer=tokenizer)
torch.save([src_outputs, trg_outputs, dec_outputs, timings], './space/data.pt')
def noisy_decode_model(model, dev_real, samples=1, alpha=1, tau=1, teacher_model=None, evaluate=True,
decoding_path=None, names=['en', 'de', 'decode'], saveall=False):
assert type(model) is FastTransformer, 'only works for fastTransformer'
logger.info('decoding from the devlopment set. beamsize={}, alpha={}, tau={}'.format(args.beam_size, args.alpha, args.temperature))
dev = dev_real
progressbar = tqdm(total=sum([1 for _ in dev]), desc='start decoding')
model.eval()
teacher_model.eval()
if decoding_path is not None:
decoding_path = decoding_path.format(args.test_set if train_real is None else 'train')
handle_dec = open(decoding_path + '.{}'.format(names[2]), 'w')
handle_src = open(decoding_path + '.{}'.format(names[0]), 'w')
handle_trg = open(decoding_path + '.{}'.format(names[1]), 'w')
# if saveall:
# handle_fer = open(decoding_path + '.{}'.format(names[3]), 'w')
corpus_size = 0
src_outputs, trg_outputs, dec_outputs, timings = [], [], [], []
all_dec_outputs = []
decoded_words, target_words, decoded_info = 0, 0, 0
attentions = None #{'source': None, 'target': None}
pad_id = model.decoder.field.vocab.stoi['<pad>']
eos_id = model.decoder.field.vocab.stoi['<eos>']
curr_time = 0
for iters, dev_batch in enumerate(dev):
start_t = time.time()
inputs, input_masks, targets, target_masks, sources, source_masks0, encoding0, batch_size = model.quick_prepare(dev_batch)
if teacher_model is not None:
encoding_teacher = teacher_model.encoding(sources, source_masks0)
batch_size, src_len, hsize = encoding0[0].size()
if samples > 1:
source_masks = source_masks0[:, None, :].expand(batch_size, samples,
src_len).contiguous().view(batch_size * samples, src_len)
encoding = [None for _ in encoding0]
for i in range(len(encoding)):
encoding[i] = encoding0[i][:, None, :].expand(
batch_size, samples, src_len, hsize).contiguous().view(batch_size * samples, src_len, hsize)
if teacher_model is not None:
for i in range(len(encoding)):
encoding_teacher[i] = encoding_teacher[i][:, None, :].expand(
batch_size, samples, src_len, hsize).contiguous().view(batch_size * samples, src_len, hsize)
def parallel():
decoder_inputs, input_reorder, decoder_masks, logits_fer = model.prepare_initial(encoding0, sources, source_masks0, input_masks,
None, None, decoding=True, mode=args.fertility_mode, N=samples, tau=tau)
if teacher_model is not None:
decoding = model(encoding, source_masks, decoder_inputs, decoder_masks, decoding=True, feedback=attentions)
student_inputs, _ = teacher_model.prepare_inputs(dev_batch, decoding, decoder_masks)
student_targets, _ = teacher_model.prepare_targets(dev_batch, decoding, decoder_masks)
out, probs = teacher_model(encoding_teacher, source_masks, student_inputs, decoder_masks, return_probs=True, decoding=False)
_, teacher_loss = model.batched_cost(student_targets, decoder_masks, probs, batched=True) # student-loss (MLE)
# reranking the translation
teacher_loss = teacher_loss.view(batch_size, samples)
decoding = decoding.view(batch_size, samples, -1)
lp = decoder_masks.sum(1).view(batch_size, samples) ** (1 - alpha)
teacher_loss = teacher_loss * Variable(lp)
return decoding, teacher_loss, input_reorder
if args.multi_run > 1:
decodings, teacher_losses, _ = zip(*[parallel() for _ in range(args.multi_run)])
maxl = max([d.size(2) for d in decodings])
decoding = Variable(sources.data.new(batch_size, samples * args.multi_run, maxl).fill_(1).long())
for i, d in enumerate(decodings):
decoding[:, i * samples: (i+1) * samples, :d.size(2)] = d
teacher_loss = torch.cat(teacher_losses, 1)
else:
decoding, teacher_loss, input_reorder = parallel()
all_dec_outputs += [(decoding.view(batch_size * samples, -1), input_reorder)]
selected_idx = (-teacher_loss).topk(1, 1)[1] # batch x 1
decoding = decoding.gather(1, selected_idx[:, :, None].expand(batch_size, 1, decoding.size(-1)))[:, 0, :]
used_t = time.time() - start_t
curr_time += used_t
real_mask = 1 - ((decoding.data == eos_id) + (decoding.data == pad_id)).float()
outputs = [model.output_decoding(d) for d in [('src', sources), ('trg', targets), ('trg', decoding)]]
corpus_size += batch_size
src_outputs += outputs[0]
trg_outputs += outputs[1]
dec_outputs += outputs[2]
timings += [used_t]
if decoding_path is not None:
for s, t, d in zip(outputs[0], outputs[1], outputs[2]):
if args.no_bpe:
s, t, d = s.replace('@@ ', ''), t.replace('@@ ', ''), d.replace('@@ ', '')
print(s, file=handle_src, flush=True)
print(t, file=handle_trg, flush=True)
print(d, file=handle_dec, flush=True)
# if saveall:
# for d, f in all_dec_outputs:
# ds = model.output_decoding(('trg', d))
# fs = model.output_decoding(('src', f))
# for dd, ff in zip(ds, fs):
# print(dd, file=handle_fer, flush=True)
# print(ff, file=handle_fer, flush=True)
progressbar.update(1)
progressbar.set_description('finishing sentences={}/batches={} speed={} sec/batch'.format(corpus_size, iters, curr_time / (1 + iters)))
if evaluate:
corpus_gleu = computeGLEU(dec_outputs, trg_outputs, corpus=True, tokenizer=tokenizer)
corpus_bleu = computeBLEU(dec_outputs, trg_outputs, corpus=True, tokenizer=tokenizer)
logger.info("The dev-set corpus GLEU = {}".format(corpus_gleu))
logger.info("The dev-set corpus BLEU = {}".format(corpus_bleu))
computeGroupBLEU(dec_outputs, trg_outputs, tokenizer=tokenizer)
torch.save([src_outputs, trg_outputs, dec_outputs, timings], './space/data.pt')
def self_improving_model(model, train, dev):
if args.tensorboard and (not args.debug):
from tensorboardX import SummaryWriter
writer = SummaryWriter('./runs/self-{}'.format(args.prefix+hp_str))
# optimizer
if args.optimizer == 'Adam':
opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], betas=(0.9, 0.98), eps=1e-9)
if args.trainable_teacher:
opt_teacher = torch.optim.Adam([p for p in teacher_model.parameters() if p.requires_grad], betas=(0.9, 0.98), eps=1e-9)
elif args.optimizer == 'RMSprop':
opt = torch.optim.RMSprop([p for p in model.parameters() if p.requires_grad], eps=1e-9)
if args.trainable_teacher:
opt_teacher = torch.optim.RMSprop([p for p in teacher_model.parameters() if p.requires_grad], eps=1e-9)
else:
raise NotImplementedError
# if resume training --
if (args.load_from is not None) and (args.resume):
with torch.cuda.device(args.gpu): # very important.
offset, opt_states = torch.load('./models/' + args.load_from + '.pt.states',
map_location=lambda storage, loc: storage.cuda())
opt.load_state_dict(opt_states)
else:
offset = 0
# metrics
best = Best(max, 'corpus_bleu', 'corpus_gleu', 'gleu', 'loss', 'i', model=model, opt=opt, path=model_name, gpu=args.gpu)
train_metrics = Metrics('train', 'loss', 'real', 'fake')
dev_metrics = Metrics('dev', 'loss', 'gleu', 'real_loss', 'fake_loss', 'distance', 'alter_loss', 'distance2', 'reordering_loss', 'corpus_gleu')
progressbar = tqdm(total=args.eval_every, desc='start training.')
# cache
samples = 100
tau = 1
caches = Cache(args.max_cache, ['src', 'trg', 'dec', 'fer'])
best_model = copy.deepcopy(model) # used for decoding
best_score = 0
# start loop
iters = offset
train = iter(train)
counters = 0
while iters <= args.maximum_steps:
iters += 1
counters += 1
batch = devol(next(train))
# prepare inputs
model.eval()
inputs, input_masks, \
targets, target_masks, \
sources, source_masks0, encoding, batch_size = model.quick_prepare(batch)
_, src_len, hsize = encoding[0].size()
trg_len = targets.size(1)
# prepare parallel -- noisy sampling
decoder_inputs, input_reorder, decoder_masks, _, pred_fer \
= model.prepare_initial(encoding, sources, source_masks0, input_masks,
None, None, decoding=True, mode='reinforce',
N=samples, tau=tau, return_samples=True)
# repeating for decoding
source_masks = source_masks0[:, None, :].expand(batch_size, samples,
src_len).contiguous().view(batch_size * samples, src_len)
for i in range(len(encoding)):
encoding[i] = encoding[i][:, None, :].expand(
batch_size, samples, src_len, hsize).contiguous().view(batch_size * samples, src_len, hsize)
# run decoding
decoding, _, probs = best_model(encoding, source_masks, decoder_inputs, decoder_masks, decoding=True, return_probs=True)
# compute GLEU score to select the best translation
trg_output = best_model.output_decoding(('trg', targets[:, None, :].expand(batch_size,
samples, trg_len).contiguous().view(batch_size * samples, trg_len)))
dec_output = best_model.output_decoding(('trg', decoding))
bleu_score = computeBLEU(dec_output, trg_output, corpus=False, tokenizer=tokenizer).contiguous().view(batch_size, samples).cuda(args.gpu)
best_index = bleu_score.max(1)[1]
def index_gather(data, index, samples):
batch_size = index.size(0)
data = data.contiguous().view(batch_size, samples, -1) # batch x samples x dim
index = index[:, None, None].expand(batch_size, 1, data.size(2))
return data.gather(1, index)[:, 0, :]
best_decoding, best_decoder_masks, best_fertilities = [index_gather(x, best_index, samples) for x in [decoding, decoder_masks, pred_fer]]
caches.add([sources, targets, best_decoding, best_fertilities],
[source_masks0, target_masks, best_decoder_masks, source_masks0],
['src', 'trg', 'dec', 'fer'])
progressbar.update(1)
progressbar.set_description('caching sentences={}/batches={}'.format(len(caches.cache), iters))
if counters == args.eval_every:
logger.info('build a new dataset from the caches')
print(len(caches.cache))
cache_data = ParallelDataset(examples=caches.cache,
fields=[('src', SRC), ('trg', TRG), ('dec', TRG), ('fer', FER)])
cache_iter = data.BucketIterator(cache_data, batch_sizes=2048, device=args.gpu, batch_size_fn=batch_size_fn)
print('done')
import sys;sys.exit(1)
if False: # iters % args.eval_every == 0:
progressbar.close()
dev_metrics.reset()
outputs_data = valid_model(model, dev, None if args.seq_dist else dev_metrics, teacher_model=None, print_out=True)
if args.tensorboard and (not args.debug):
writer.add_scalar('dev/GLEU_sentence_', dev_metrics.gleu, iters)
writer.add_scalar('dev/Loss', dev_metrics.loss, iters)
writer.add_scalar('dev/GLEU_corpus_', outputs_data['corpus_gleu'], iters)
writer.add_scalar('dev/BLEU_corpus_', outputs_data['corpus_bleu'], iters)
if not args.debug:
best.accumulate(outputs_data['corpus_bleu'], outputs_data['corpus_gleu'], dev_metrics.gleu, dev_metrics.loss, iters)
logger.info('the best model is achieved at {}, average greedy GLEU={}, corpus GLEU={}, corpus BLEU={}'.format(
best.i, best.gleu, best.corpus_gleu, best.corpus_bleu))
logger.info('model:' + args.prefix + hp_str)
# ---set-up a new progressor---
progressbar = tqdm(total=args.eval_every, desc='start training.')
if args.mode == 'train':
logger.info('starting training')
train_model(model, train_real, dev_real, teacher_model)
elif args.mode == 'self':
logger.info('starting self-training')
self_improving_model(model, train_real, dev_real)
elif args.mode == 'test':
logger.info('starting decoding from the pre-trained model, test...')
names = ['dev.src.b{}={}.{}'.format(args.beam_size, args.load_from, args,fertility_mode),
'dev.trg.b{}={}.{}'.format(args.beam_size, args.load_from, args,fertility_mode),
'dev.dec.b{}={}.{}'.format(args.beam_size, args.load_from, args,fertility_mode)]
decode_model(model, None, dev_real, evaluate=True, decoding_path=decoding_path if not args.no_write else None, names=names)
elif args.mode == 'test_noisy':
logger.info('starting decoding from the pre-trained model, test...')
names = ['dev.src.b{}={}.noise{}'.format(args.beam_size, args.load_from, args.beam_size),
'dev.trg.b{}={}.noise{}'.format(args.beam_size, args.load_from, args.beam_size),
'dev.dec.b{}={}.noise{}'.format(args.beam_size, args.load_from, args.beam_size),
'dev.fer.b{}={}.noise{}'.format(args.beam_size, args.load_from, args.beam_size)]
noisy_decode_model(model, dev_real, samples=args.beam_size, alpha=args.alpha, tau=args.temperature,
teacher_model=teacher_model, evaluate=True, decoding_path=decoding_path if not args.no_write else None,
names=names, saveall=True)
else:
logger.info('starting decoding from the pre-trained model, build the course dataset...')
names = ['src.b{}'.format(args.beam_size), 'trg.b{}'.format(args.beam_size), 'dec.b{}'.format(args.beam_size)]
decode_model(model, train_real, dev_real, decoding_path=decoding_path if not args.no_write else None, names=names)
logger.info("done.")
