#

ゲーテッドリニアユニットとバリアント

これにより、自動回帰用の単純なトランスフォーマーモデルがトレーニングされます。位置ごとのフィードフォワードネットワークにはさまざまなバリエーションを試します。再利用可能および設定可能なものはで定義されています

。configs.py

16import torch
17from labml import experiment
18from labml.configs import option
19from labml.utils.pytorch import get_modules
20from labml_helpers.module import Module
21
22from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
23from labml_nn.transformers import Encoder, Generator, TransformerConfigs
24from labml_nn.transformers.utils import subsequent_mask

#

自動回帰モデル

27class AutoregressiveModel(Module):

#

32    def __init__(self, src_embed: Module, encoder: Encoder, generator: Generator):
33        super().__init__()

#

トークン埋め込みモジュール

35        self.src_embed = src_embed

#

トランスベースのエンコーダ

37        self.encoder = encoder

#

次のトークン生成レイヤー。これにより、次のトークンのロジットが返されます

40        self.generator = generator

#

これは最初の呼び出しで初期化されます。

42        self.src_mask = None

#

44    def forward(self, src: torch.Tensor):

#

次のマスクを作成して、トランスフォーマーが過去のトークンにしか注目できないようにします。

46        if self.src_mask is None or self.src_mask.size(0) != len(src):
47            self.src_mask = subsequent_mask(len(src)).to(src.device)

#

トークン (src ) を埋め込み、トランスフォーマーに通します

49        res = self.encoder(self.src_embed(src), self.src_mask)

#

次のトークンのロジットを生成

51        return self.generator(res), None

#

コンフィギュレーション

デフォルトの設定は、実験を開始したときに上書きでき、また上書きされます。

54class Configs(NLPAutoRegressionConfigs):

#

61    transformer: TransformerConfigs
62    model: AutoregressiveModel

#

自己回帰モデルを初期化

65@option(Configs.model)
66def autoregressive_model(c: Configs):

#

70    m = AutoregressiveModel(c.transformer.src_embed, c.transformer.encoder, c.transformer.generator)
71    return m.to(c.device)

#

自己回帰モデルの設定可能なトランスエンコーダーを初期化します。

74@option(Configs.transformer)
75def transformer_c(c: Configs):

#

79    tc = TransformerConfigs()
80    tc.n_src_vocab = c.n_tokens
81    tc.n_tgt_vocab = c.n_tokens
82
83    return tc

#

86def main():

#

実験を作成

88    experiment.create(name="glu_variants")

#

コンフィグの作成

90    conf = Configs()

#

構成をロード

92    experiment.configs(conf,

#

オーバーライドする設定の辞書

94                       {'tokenizer': 'character',
95                        'prompt_separator': '',
96                        'prompt': 'It is ',
97                        'text': 'tiny_shakespeare',
98
99                        'optimizer.optimizer': 'Noam',
100                        'optimizer.learning_rate': 1.,
101                        'optimizer.d_model': 256,
102
103                        'seq_len': 1024,
104                        'epochs': 128,
105                        'batch_size': 6,
106                        'inner_iterations': 10,

#

GLU バリアント、GLU、バイリニア、RegLU、GEGLU、SwiGLU のいずれか

これらは設定可能な FFN 実装で定義されています

112                        'transformer.ffn.glu_variant': 'Bilinear',

#

変圧器構成

115                        'transformer.d_model': 256,
116                        'transformer.ffn.d_ff': 1024,
117                        'transformer.n_heads': 8,
118                        'transformer.n_layers': 6})

#

これはモデルを初期化するために必要です

121    conf.n_tokens = conf.text.n_tokens

#

保存および読み込み用のモデルを設定する

124    experiment.add_pytorch_models(get_modules(conf))

#

実験を始める

127    with experiment.start():

#

TrainValidConfigs.run

129        conf.run()
130
131
132if __name__ == '__main__':
133    main()