Simple Neural Machine Translation (Simple-NMT)

Features

Implemented Optimization Algorithms

Maximum Likelihood Estimation (MLE)

Minimum Risk Training (MRT)

Dual Supervised Learning (DSL)

Requirements

Evaluation

Results

Setup

Usage

Training

>> python train.py -h
usage: train.py [-h] --model_fn MODEL_FN --train TRAIN --valid VALID --lang
 LANG [--gpu_id GPU_ID] [--off_autocast]
 [--batch_size BATCH_SIZE] [--n_epochs N_EPOCHS]
 [--verbose VERBOSE] [--init_epoch INIT_EPOCH]
 [--max_length MAX_LENGTH] [--dropout DROPOUT]
 [--word_vec_size WORD_VEC_SIZE] [--hidden_size HIDDEN_SIZE]
 [--n_layers N_LAYERS] [--max_grad_norm MAX_GRAD_NORM]
 [--iteration_per_update ITERATION_PER_UPDATE] [--lr LR]
 [--lr_step LR_STEP] [--lr_gamma LR_GAMMA]
 [--lr_decay_start LR_DECAY_START] [--use_adam] [--use_radam]
 [--rl_lr RL_LR] [--rl_n_samples RL_N_SAMPLES]
 [--rl_n_epochs RL_N_EPOCHS] [--rl_n_gram RL_N_GRAM]
 [--rl_reward RL_REWARD] [--use_transformer]
 [--n_splits N_SPLITS]
optional arguments:
 -h, --help show this help message and exit
 --model_fn MODEL_FN Model file name to save. Additional information would
 be annotated to the file name.
 --train TRAIN Training set file name except the extention. (ex:
 train.en --> train)
 --valid VALID Validation set file name except the extention. (ex:
 valid.en --> valid)
 --lang LANG Set of extention represents language pair. (ex: en +
 ko --> enko)
 --gpu_id GPU_ID GPU ID to train. Currently, GPU parallel is not
 supported. -1 for CPU. Default=-1
 --off_autocast Turn-off Automatic Mixed Precision (AMP), which speed-
 up training.
 --batch_size BATCH_SIZE
 Mini batch size for gradient descent. Default=32
 --n_epochs N_EPOCHS Number of epochs to train. Default=20
 --verbose VERBOSE VERBOSE_SILENT, VERBOSE_EPOCH_WISE, VERBOSE_BATCH_WISE
 = 0, 1, 2. Default=2
 --init_epoch INIT_EPOCH
 Set initial epoch number, which can be useful in
 continue training. Default=1
 --max_length MAX_LENGTH
 Maximum length of the training sequence. Default=100
 --dropout DROPOUT Dropout rate. Default=0.2
 --word_vec_size WORD_VEC_SIZE
 Word embedding vector dimension. Default=512
 --hidden_size HIDDEN_SIZE
 Hidden size of LSTM. Default=768
 --n_layers N_LAYERS Number of layers in LSTM. Default=4
 --max_grad_norm MAX_GRAD_NORM
 Threshold for gradient clipping. Default=5.0
 --iteration_per_update ITERATION_PER_UPDATE
 Number of feed-forward iterations for one parameter
 update. Default=1
 --lr LR Initial learning rate. Default=1.0
 --lr_step LR_STEP Number of epochs for each learning rate decay.
 Default=1
 --lr_gamma LR_GAMMA Learning rate decay rate. Default=0.5
 --lr_decay_start LR_DECAY_START
 Learning rate decay start at. Default=10
 --use_adam Use Adam as optimizer instead of SGD. Other lr
 arguments should be changed.
 --use_radam Use rectified Adam as optimizer. Other lr arguments
 should be changed.
 --rl_lr RL_LR Learning rate for reinforcement learning. Default=0.01
 --rl_n_samples RL_N_SAMPLES
 Number of samples to get baseline. Default=1
 --rl_n_epochs RL_N_EPOCHS
 Number of epochs for reinforcement learning.
 Default=10
 --rl_n_gram RL_N_GRAM
 Maximum number of tokens to calculate BLEU for
 reinforcement learning. Default=6
 --rl_reward RL_REWARD
 Method name to use as reward function for RL training.
 Default=gleu
 --use_transformer Set model architecture as Transformer.
 --n_splits N_SPLITS Number of heads in multi-head attention in
 Transformer. Default=8

Seq2Seq

>> python train.py --train ./data/corpus.shuf.train.tok.bpe --valid ./data/corpus.shuf.valid.tok.bpe --lang enko \
--gpu_id 0 --batch_size 128 --n_epochs 30 --max_length 100 --dropout .2 \
--word_vec_size 512 --hidden_size 768 --n_layers 4 --max_grad_norm 1e+8 --iteration_per_update 2 \
--lr 1e-3 --lr_step 0 --use_adam --rl_n_epochs 0 \
--model_fn ./model.pth

To continue with RL training

>> python continue_train.py --load_fn ./model.pth --model_fn ./model.rl.pth \
--init_epoch 31 --iteration_per_update 1 --max_grad_norm 5

Transformer

>> python train.py --train ./data/corpus.shuf.train.tok.bpe --valid ./data/corpus.shuf.valid.tok.bpe --lang enko \
--gpu_id 0 --batch_size 128 --n_epochs 30 --max_length 100 --dropout .2 \
--hidden_size 768 --n_layers 4 --max_grad_norm 1e+8 --iteration_per_update 32 \
--lr 1e-3 --lr_step 0 --use_adam --use_transformer --rl_n_epochs 0 \
--model_fn ./model.pth

Dual Supervised Learning

>> python lm_train.py --train ./data/corpus.shuf.train.tok.bpe --valid ./data/corpus.shuf.valid.tok.bpe --lang enko \
--gpu_id 0 --batch_size 256 --n_epochs 20 --max_length 64 --dropout .2 \
--word_vec_size 512 --hidden_size 768 --n_layers 4 --max_grad_norm 1e+8 \
--model_fn ./lm.pth

>> python dual_train.py --train ./data/corpus.shuf.train.tok.bpe --valid ./data/corpus.shuf.valid.tok.bpe --lang enko \
--gpu_id 0 --batch_size 64 --n_epochs 40 --max_length 64 --dropout .2 \
--word_vec_size 512 --hidden_size 768 --n_layers 4 --max_grad_norm 1e+8 --iteration_per_update 4 \
--dsl_n_warmup_epochs 30 --dsl_lambda 1e-2 \
--lm_fn ./lm.pth \
--model_fn ./model.pth

Inference

>> python translate.py -h
usage: translate.py [-h] --model_fn MODEL_FN [--gpu_id GPU_ID]
 [--batch_size BATCH_SIZE] [--max_length MAX_LENGTH]
 [--n_best N_BEST] [--beam_size BEAM_SIZE] [--lang LANG]
 [--length_penalty LENGTH_PENALTY]
optional arguments:
 -h, --help show this help message and exit
 --model_fn MODEL_FN Model file name to use
 --gpu_id GPU_ID GPU ID to use. -1 for CPU. Default=-1
 --batch_size BATCH_SIZE
 Mini batch size for parallel inference. Default=128
 --max_length MAX_LENGTH
 Maximum sequence length for inference. Default=255
 --n_best N_BEST Number of best inference result per sample. Default=1
 --beam_size BEAM_SIZE
 Beam size for beam search. Default=5
 --lang LANG Source language and target language. Example: enko
 --length_penalty LENGTH_PENALTY
 Length penalty parameter that higher value produce
 shorter results. Default=1.2

>> python translate.py --model_fn ./model.pth --gpu_id 0 --lang enko < test.txt > test.result.txt

Translation Examples

References