#

トランスフォーマー自動回帰実験

これは、「必要なのは注意だけ」で紹介したシンプルなトランスフォーマーをNLP自動回帰タスク（Tiny Shakespeareデータセットを使用）でトレーニングします。

16import torch
17from torch import nn
18
19from labml import experiment
20from labml.configs import option
21from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
22from labml_nn.transformers import TransformerConfigs, Encoder
23from labml_nn.transformers.utils import subsequent_mask

#

自己回帰モデル

26class AutoregressiveTransformer(nn.Module):

#

encoder 変圧器エンコーダです
src_embed はトークン埋め込みモジュールです (位置エンコーディング付き)

generator ロジットを生成する最後の完全接続層です。

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

#

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

#

マスクは最初の呼び出しで初期化されます

43        self.mask = None

#

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

#

マスクが初期化されていない場合やマスクのサイズが異なる場合は、後続のマスクを作成します

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

#

次にマスクすると、トークンがマスクされ、将来のトークンが見えなくなります

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

#

位置エンコーディングによるトークンの埋め込みを取得

52        x = self.src_embed(x)

#

トランスエンコーダー

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

#

ロジットを取得

56        x = self.generator(x)

#

結果を返します（トレーナーはRNNでも使用されるため、2番目の値は状態用です）

60        return x, None

#

コンフィギュレーション

これは以下から継承されます NLPAutoRegressionConfigs

63class Configs(NLPAutoRegressionConfigs):

#

GPT モデル

72    model: AutoregressiveTransformer

#

変圧器

74    transformer: TransformerConfigs

#

変圧器構成

77@option(Configs.transformer, 'Transformer')
78def _transformer_configs(c: Configs):

#

設定可能なトランス実装を使用しています

85    conf = TransformerConfigs()

#

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

87    conf.n_src_vocab = c.n_tokens
88    conf.n_tgt_vocab = c.n_tokens

#

90    conf.d_model = c.d_model

#

93    return conf

#

GPT モデルの作成と重みの初期化

96@option(Configs.model)
97def _model(c: Configs):

#

101    m = AutoregressiveTransformer(c.transformer.encoder,
102                                  c.transformer.src_embed,
103                                  c.transformer.generator).to(c.device)
104
105    return m

#

108def main():

#

実験を作成

110    experiment.create(name="transformer")

#

コンフィグの作成

112    conf = Configs()

#

オーバーライド設定

114    experiment.configs(conf, {

#

キャラクターレベルのトークナイザーを使う

116        'tokenizer': 'character',

#

プロンプトセパレータが空白

118        'prompt_separator': '',

#

サンプリングの開始プロンプト

120        'prompt': 'It is ',

#

タイニー・シェイクスピア・データセットを使う

122        'text': 'tiny_shakespeare',

#

コンテキストサイズを次の値にしてください $256$

125        'seq_len': 512,

#

32 エポックのトレーニング

127        'epochs': 32,

#

バッチサイズ $32$

129        'batch_size': 16,

#

エポックごとにトレーニングと検証を切り替える $10$

132        'inner_iterations': 10,

#

モデルサイズ

135        'd_model': 256,
136        'transformer.n_heads': 16,
137        'transformer.ffn.d_ff': 1024,

#

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