#

トレイン・フィードバック・トランス

これにより、自己回帰用のフィードバックトランスフォーマーモデルがトレーニングされます。オリジナルのフィードバックトランスフォーマーか、キーと値が事前に計算された新しいバージョンを選択できます

。

これは、Tiny ShakespeareデータセットのフィードバックトランスフォーマーをトレーニングするためのColabノートブックです。

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

#

自動回帰モデル

31class AutoregressiveModel(Module):

#

36    def __init__(self, n_vocab: int, d_model: int, transformer: Module):
37        super().__init__()

#

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

39        self.src_embed = nn.Embedding(n_vocab, d_model)
40        self.transformer = transformer
41        self.generator = nn.Linear(d_model, n_vocab)

#

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

#

トークンを埋め込む

45        x = self.src_embed(x)

#

変圧器に通してください

47        res = self.transformer(x)

#

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

49        return self.generator(res), None

#

コンフィギュレーション

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

52class Configs(NLPAutoRegressionConfigs):

#

59    model: AutoregressiveModel
60
61    d_model: int = 512
62    heads: int = 8
63    dropout: float = 0.0
64    d_ff: int = 2048
65    n_layers: int = 6

#

オリジナルのフィードバックトランスフォーマーを作成

68@option(Configs.model)
69def feedback_transformer(c: Configs):

#

73    from labml_nn.transformers.feedback import FeedbackTransformer, FeedbackTransformerLayer, \
74        FeedbackAttention, FeedForward
75
76    return AutoregressiveModel(
77        c.n_tokens, c.d_model,
78        FeedbackTransformer(
79            FeedbackTransformerLayer(d_model=c.d_model,
80                                     attn=FeedbackAttention(c.heads, c.d_model, c.dropout),
81                                     feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout),
82                                     dropout_prob=c.dropout),
83            c.n_layers)).to(c.device)

#

事前に計算されたキーと値を使用して、更新されたフィードバックトランスフォーマーを作成します。

86@option(Configs.model)
87def feedback_transformer_kv(c: Configs):

#

91    from labml_nn.transformers.feedback import FeedbackTransformerKV, FeedbackTransformerLayer, \
92        FeedbackAttention, FeedForward
93
94    return AutoregressiveModel(
95        c.n_tokens, c.d_model,
96        FeedbackTransformerKV(
97            FeedbackTransformerLayer(d_model=c.d_model,
98                                     attn=FeedbackAttention(c.heads, c.d_model, c.dropout,
99                                                            is_kv_precomputed=True),
100                                     feed_forward=FeedForward(c.d_model, c.d_ff, c.dropout),
101                                     dropout_prob=c.dropout),
102            c.n_layers, c.d_model, c.heads)).to(c.device)

#

105def main():

#

実験を作成

107    experiment.create(name="feedback_transformer")

#

コンフィグの作成

109    conf = Configs()

#

構成をロード

111    experiment.configs(conf,

#

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

113                       {'tokenizer': 'character',
114                        'text': 'tiny_shakespeare',
115                        'optimizer.learning_rate': 1.0,
116                        'optimizer.optimizer': 'Noam',
117                        'prompt': 'It is',
118                        'prompt_separator': '',

#

feedback_transformer オリジナルのフィードバックトランスに使用

121                        'model': 'feedback_transformer_kv',
122
123                        'train_loader': 'shuffled_train_loader',
124                        'valid_loader': 'shuffled_valid_loader',
125
126                        'seq_len': 128,
127                        'epochs': 128,
128                        'batch_size': 64,
129                        'inner_iterations': 25})

#

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

132    experiment.add_pytorch_models(get_modules(conf))

#

実験を始める

135    with experiment.start():

#

トレーニングループを実行する

137        conf.run()
138
139
140if __name__ == '__main__':
141    main()