#

注意 mLP (gmLP) 实验

这是一项带注释的 PyTorch 实验，用于训练 gmLP 模型。本文还应用了随机深度正则化，在训练过程中会随机删除一些图层。我们没有在这里实现这一点。

16from labml import experiment
17from labml.configs import option
18from labml_nn.transformers import TransformerConfigs
19from labml_nn.transformers.basic.autoregressive_experiment import Configs as BasicAutoRegressionConfigs
20from labml_nn.transformers.gmlp import GMLPBlock

#

配置

这继承自训练循环和简单变压器自回归自然语言处理任务的配置。

23class Configs(BasicAutoRegressionConfigs):

#

变压器

32    transformer: TransformerConfigs = 'gMLP'

#

gmLP Block

34    gmlp: GMLPBlock

#

d_ffn 用于 gMLP 投影层

36    d_ffn: int = 2048

#

创建一个 GmLP 区块

39@option(Configs.gmlp, 'gMLP')
40def _gmlp_configs(c: Configs):

#

44    return GMLPBlock(c.d_model, c.d_ffn, c.seq_len)

#

变压器配置

47@option(Configs.transformer, 'gMLP')
48def _transformer_configs(c: Configs):

#

我们使用我们的可配置变压器实现

55    conf = TransformerConfigs()

#

设置嵌入和生成 logit 的词汇量大小

57    conf.n_src_vocab = c.n_tokens
58    conf.n_tgt_vocab = c.n_tokens

#

设置模型尺寸

60    conf.d_model = c.d_model

#

将编码器层替换为 GmLP 层

62    conf.encoder_layer = c.gmlp
63
64    return conf

#

67def main():

#

创建实验

69    experiment.create(name="gMLP")

#

创建配置

71    conf = Configs()

#

覆盖配置

73    experiment.configs(conf, {

#

使用角色等级分词器

75        'tokenizer': 'character',

#

提示分隔符为空

77        'prompt_separator': '',

#

开始采样提示

79        'prompt': 'It is ',

#

使用小莎士比亚数据集

81        'text': 'tiny_shakespeare',

#

使用上下文大小为 $256$

84        'seq_len': 256,

#

为 $128$ 时代而训练

86        'epochs': 128,

#

批量大小 $32$

88        'batch_size': 32,

#

在训练和验证之间切换每个纪元的 $10$ 次数

91        'inner_iterations': 10,

#

型号尺寸

94        'd_model': 512,
95        'd_ffn': 2048,

#

使用 Noam 优化器

98        'optimizer.optimizer': 'Noam',
99        'optimizer.learning_rate': 1.,
100    })

#

设置用于保存和加载的模型

103    experiment.add_pytorch_models({'model': conf.model})

#

开始实验

106    with experiment.start():

#

跑步训练

108        conf.run()

#

112if __name__ == '__main__':
113    main()