#

アテンションフリー変圧器（AFT）実験

これは、一般的なトレーニングループと自己回帰型NLPタスクの設定に基づいています。

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 ロジットを生成する最後の完全接続層です。

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でも使用されるため、2番目の値は状態用です）

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()

#

埋め込みやロジットの生成に使用するボキャブラリーサイズを設定

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()