#

无注意变形金刚 (AFT) 实验

这是一项带注释的 PyTorch 实验，用于训练 A FT 模型。

14import torch
15
16from labml import experiment
17from labml.configs import option
18from labml_helpers.module import Module
19from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
20from labml_nn.transformers import TransformerConfigs, Encoder
21from labml_nn.transformers.utils import subsequent_mask

#

简单的自回归模型

这包括令牌嵌入层、变压器编码器和给出令牌日志的最终线性层。

24class AutoregressiveTransformer(Module):

#

encoder 是变压器编码器
src_embed 是令牌嵌入模块（带有位置编码）
generator 是给出 logit 的最后一个完全连接的层。

32    def __init__(self, encoder: Encoder, src_embed: Module, generator: Module):

#

39        super().__init__()
40        self.src_embed = src_embed
41        self.encoder = encoder
42        self.generator = generator

#

掩码将在第一次调用时初始化

45        self.mask = None

#

47    def forward(self, x: torch.Tensor):

#

如果掩码未初始化或掩码大小不同，则创建后续掩码

50        if self.mask is None or self.mask.size(0) != len(x):

#

后续的掩码，将掩盖令牌以免看到未来的代币

52            self.mask = subsequent_mask(len(x)).to(x.device)

#

使用位置编码获取令牌嵌入

55        x = self.src_embed(x)

#

变压器编码

57        x = self.encoder(x, self.mask)

#

获取日志

59        x = self.generator(x)

#

返回结果（第二个值用于状态，因为我们的训练器也与 RNN 一起使用）

63        return x, None

#

配置

这继承自 NLPAutoRegressionConfigs

66class Configs(NLPAutoRegressionConfigs):

#

GPT 型号

75    model: AutoregressiveTransformer

#

变压器

77    transformer: TransformerConfigs
78
79    local_window_size: int = 32

#

变压器配置

82@option(Configs.transformer, 'Transformer')
83def _transformer_configs(c: Configs):

#

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

90    conf = TransformerConfigs()

#

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

92    conf.n_src_vocab = c.n_tokens
93    conf.n_tgt_vocab = c.n_tokens

#

设置嵌入大小

95    conf.d_model = c.d_model

#

用 AFT 本地模块替换自我注意

97    from labml_nn.transformers.aft import AFTLocal
98    conf.encoder_attn = AFTLocal(c.d_model, c.seq_len, c.local_window_size)

#

101    return conf

#

创建自动回归模型

104@option(Configs.model)
105def _model(c: Configs):

#

109    m = AutoregressiveTransformer(c.transformer.encoder,
110                                  c.transformer.src_embed,
111                                  c.transformer.generator).to(c.device)
112
113    return m

#

116def main():

#

创建实验

118    experiment.create(name="aft")

#

创建配置

120    conf = Configs()

#

覆盖配置

122    experiment.configs(conf, {

#

使用角色等级分词器

124        'tokenizer': 'character',

#

提示分隔符为空

126        'prompt_separator': '',

#

开始采样提示

128        'prompt': 'It is ',

#

使用小莎士比亚数据集

130        'text': 'tiny_shakespeare',

#

使用上下文大小为 $128$

133        'seq_len': 256,

#

为 $32$ 时代而训练

135        'epochs': 128,

#

批量大小 $128$

137        'batch_size': 32,

#

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

140        'inner_iterations': 10,

#

嵌入大小

143        'd_model': 128,

#

FFN 隐藏尺寸大小

145        'transformer.ffn.d_ff': 256,

#

优化器

148        'optimizer.optimizer': 'Noam',
149        'optimizer.learning_rate': 1.,
150    })

#

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

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

#

开始实验

156    with experiment.start():

#

跑步训练

158        conf.run()

#

162if __name__ == '__main__':
163    main()