config.py 3.2 KB
Newer Older
Y
ying 已提交
1 2 3 4 5
class TrainTaskConfig(object):
    use_gpu = False
    # the epoch number to train.
    pass_num = 2

6
    # the number of sequences contained in a mini-batch.
Y
ying 已提交
7 8
    batch_size = 64

9
    # the hyper parameters for Adam optimizer.
Y
ying 已提交
10 11 12 13 14
    learning_rate = 0.001
    beta1 = 0.9
    beta2 = 0.98
    eps = 1e-9

15
    # the parameters for learning rate scheduling.
16 17
    warmup_steps = 4000

18 19
    # the directory for saving trained models.
    model_dir = "trained_models"
20 21 22 23


class InferTaskConfig(object):
    use_gpu = False
24 25
    # the number of examples in one run for sequence generation.
    # currently the batch size can only be set to 1.
26 27
    batch_size = 1

28
    # the parameters for beam search.
29 30
    beam_size = 5
    max_length = 30
31
    # the number of decoded sentences to output.
32 33
    n_best = 1

34 35
    # the directory for loading the trained model.
    model_path = "trained_models/pass_1.infer.model"
36

Y
ying 已提交
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54

class ModelHyperParams(object):
    # Dictionary size for source and target language. This model directly uses
    # paddle.dataset.wmt16 in which <bos>, <eos> and <unk> token has
    # alreay been added, but the <pad> token is not added. Transformer requires
    # sequences in a mini-batch are padded to have the same length. A <pad> token is
    # added into the original dictionary in paddle.dateset.wmt16.

    # size of source word dictionary.
    src_vocab_size = 10000
    # index for <pad> token in source language.
    src_pad_idx = src_vocab_size

    # size of target word dictionay
    trg_vocab_size = 10000
    # index for <pad> token in target language.
    trg_pad_idx = trg_vocab_size

55 56 57 58 59
    # index for <bos> token
    bos_idx = 0
    # index for <eos> token
    eos_idx = 1

Y
ying 已提交
60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90
    # position value corresponding to the <pad> token.
    pos_pad_idx = 0

    # max length of sequences. It should plus 1 to include position
    # padding token for position encoding.
    max_length = 50

    # the dimension for word embeddings, which is also the last dimension of
    # the input and output of multi-head attention, position-wise feed-forward
    # networks, encoder and decoder.

    d_model = 512
    # size of the hidden layer in position-wise feed-forward networks.
    d_inner_hid = 1024
    # the dimension that keys are projected to for dot-product attention.
    d_key = 64
    # the dimension that values are projected to for dot-product attention.
    d_value = 64
    # number of head used in multi-head attention.
    n_head = 8
    # number of sub-layers to be stacked in the encoder and decoder.
    n_layer = 6
    # dropout rate used by all dropout layers.
    dropout = 0.1


# Names of position encoding table which will be initialized externally.
pos_enc_param_names = (
    "src_pos_enc_table",
    "trg_pos_enc_table", )

91 92
# Names of all data layers in encoder listed in order.
encoder_input_data_names = (
Y
ying 已提交
93 94
    "src_word",
    "src_pos",
95 96
    "src_slf_attn_bias",
    "src_data_shape", )
97 98 99

# Names of all data layers in decoder listed in order.
decoder_input_data_names = (
Y
ying 已提交
100 101 102 103
    "trg_word",
    "trg_pos",
    "trg_slf_attn_bias",
    "trg_src_attn_bias",
104
    "trg_data_shape",
105 106 107 108
    "enc_output", )

# Names of label related data layers listed in order.
label_data_names = (
109 110
    "lbl_word",
    "lbl_weight", )