#

FNet 实验

这是一个带注释的 PyTorch 实验，用于训练 FNet 模型。

15import torch
16from torch import nn
17
18from labml import experiment
19from labml.configs import option
20from labml_helpers.module import Module
21from labml_nn.experiments.nlp_classification import NLPClassificationConfigs
22from labml_nn.transformers import Encoder
23from labml_nn.transformers import TransformerConfigs

#

基于变压器的分类器模型

26class TransformerClassifier(nn.Module):

#

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

30    def __init__(self, encoder: Encoder, src_embed: Module, generator: nn.Linear):

#

37        super().__init__()
38        self.src_embed = src_embed
39        self.encoder = encoder
40        self.generator = generator

#

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

#

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

44        x = self.src_embed(x)

#

变压器编码

46        x = self.encoder(x, None)

#

获取分类日志。

我们将令[CLS] 牌设置在序列的最后一个位置。这是由，wherx[-1] ex 是形状提取

的[seq_len, batch_size, d_model]

52        x = self.generator(x[-1])

#

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

56        return x, None

#

配置

这继承自 NLPClassificationConfigs

59class Configs(NLPClassificationConfigs):

#

分类模型

68    model: TransformerClassifier

#

变压器

70    transformer: TransformerConfigs

#

变压器配置

73@option(Configs.transformer)
74def _transformer_configs(c: Configs):

#

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

81    conf = TransformerConfigs()

#

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

83    conf.n_src_vocab = c.n_tokens
84    conf.n_tgt_vocab = c.n_tokens

#

87    return conf

#

创建可以取代变压器编码器层中的自我注意力的FNetMix 模块。

90@option(TransformerConfigs.encoder_attn)
91def fnet_mix():

#

97    from labml_nn.transformers.fnet import FNetMix
98    return FNetMix()

#

创建分类模型

101@option(Configs.model)
102def _model(c: Configs):

#

106    m = TransformerClassifier(c.transformer.encoder,
107                              c.transformer.src_embed,
108                              nn.Linear(c.d_model, c.n_classes)).to(c.device)
109
110    return m

#

113def main():

#

创建实验

115    experiment.create(name="fnet")

#

创建配置

117    conf = Configs()

#

覆盖配置

119    experiment.configs(conf, {

#

使用世界级分词器

121        'tokenizer': 'basic_english',

#

为 $32$ 时代而训练

124        'epochs': 32,

#

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

127        'inner_iterations': 10,

#

变压器配置（与默认值相同）

130        'transformer.d_model': 512,
131        'transformer.ffn.d_ff': 2048,
132        'transformer.n_heads': 8,
133        'transformer.n_layers': 6,

#

使用 FNet 而不是自我关注

137        'transformer.encoder_attn': 'fnet_mix',

#

使用 Noam 优化器

140        'optimizer.optimizer': 'Noam',
141        'optimizer.learning_rate': 1.,
142    })

#

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

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

#

开始实验

148    with experiment.start():

#

跑步训练

150        conf.run()

#

154if __name__ == '__main__':
155    main()