Seq2Seq in PyTorch
This is a complete suite for training sequence-to-sequence models in PyTorch. It consists of several models and code to both train and infer using them.
Using this code you can train:
- Neural-machine-translation (NMT) models
- Language models
- Image to caption generation
- Skip-thought sentence representations
- And more...
Installation
git clone --recursive https://github.com/eladhoffer/seq2seq.pytorch
cd seq2seq.pytorch; python setup.py develop
Models
Models currently available:
- Simple Seq2Seq recurrent model
- Recurrent Seq2Seq with attentional decoder
- Google neural machine translation (GNMT) recurrent model
- Transformer - attention-only model from "Attention Is All You Need"
Datasets
Datasets currently available:
- WMT16
- WMT17
- OpenSubtitles 2016
- COCO image captions
- Conceptual captions
All datasets can be tokenized using 3 available segmentation methods:
- Character based segmentation
- Word based segmentation
- Byte-pair-encoding (BPE) as suggested by bpe with selectable number of tokens.
After choosing a tokenization method, a vocabulary will be generated and saved for future inference.
Training methods
The models can be trained using several methods:
- Basic Seq2Seq - given encoded sequence, generate (decode) output sequence. Training is done with teacher-forcing.
- Multi Seq2Seq - where several tasks (such as multiple languages) are trained simultaneously by using the data sequences as both input to the encoder and output for decoder.
- Image2Seq - used to train image to caption generators.
Usage
Example training scripts are available in scripts folder. Inference examples are available in examples folder.
- example for training a transformer on WMT16 according to original paper regime:
DATASET=${1:-"WMT16_de_en"}
DATASET_DIR=${2:-"./data/wmt16_de_en"}
OUTPUT_DIR=${3:-"./results"}
WARMUP="4000"
LR0="512**(-0.5)"
python main.py \
--save transformer \
--dataset ${DATASET} \
--dataset-dir ${DATASET_DIR} \
--results-dir ${OUTPUT_DIR} \
--model Transformer \
--model-config "{'num_layers': 6, 'hidden_size': 512, 'num_heads': 8, 'inner_linear': 2048}" \
--data-config "{'moses_pretok': True, 'tokenization':'bpe', 'num_symbols':32000, 'shared_vocab':True}" \
--b 128 \
--max-length 100 \
--device-ids 0 \
--label-smoothing 0.1 \
--trainer Seq2SeqTrainer \
--optimization-config "[{'step_lambda':
\"lambda t: { \
'optimizer': 'Adam', \
'lr': ${LR0} * min(t ** -0.5, t * ${WARMUP} ** -1.5), \
'betas': (0.9, 0.98), 'eps':1e-9}\"
}]"
- example for training attentional LSTM based model with 3 layers in both encoder and decoder:
python main.py \
--save de_en_wmt17 \
--dataset ${DATASET} \
--dataset-dir ${DATASET_DIR} \
--results-dir ${OUTPUT_DIR} \
--model RecurrentAttentionSeq2Seq \
--model-config "{'hidden_size': 512, 'dropout': 0.2, \
'tie_embedding': True, 'transfer_hidden': False, \
'encoder': {'num_layers': 3, 'bidirectional': True, 'num_bidirectional': 1, 'context_transform': 512}, \
'decoder': {'num_layers': 3, 'concat_attention': True,\
'attention': {'mode': 'dot_prod', 'dropout': 0, 'output_transform': True, 'output_nonlinearity': 'relu'}}}" \
--data-config "{'moses_pretok': True, 'tokenization':'bpe', 'num_symbols':32000, 'shared_vocab':True}" \
--b 128 \
--max-length 80 \
--device-ids 0 \
--trainer Seq2SeqTrainer \
--optimization-config "[{'epoch': 0, 'optimizer': 'Adam', 'lr': 1e-3},
{'epoch': 6, 'lr': 5e-4},
{'epoch': 8, 'lr':1e-4},
{'epoch': 10, 'lr': 5e-5},
{'epoch': 12, 'lr': 1e-5}]" \
154 Jan 01, 2023
537 Jan 05, 2023
100 Dec 12, 2022
61 Dec 13, 2022
50 Mar 22, 2022
10 Jan 06, 2023
1.1k Dec 25, 2022
14 Dec 15, 2021
409 Jan 06, 2023
27 Oct 30, 2022
133 Sep 20, 2022
20 Dec 14, 2022
151 Jan 05, 2023
6 Sep 02, 2022
122 Nov 17, 2022
877 Jan 05, 2023
15 Dec 11, 2022
64 Aug 17, 2022
6 Feb 10, 2022
26 Dec 05, 2022