#

ディープ・ノーム・エクスペリメント

13import copy
14
15import torch
16import torch.nn as nn
17
18from labml import experiment
19from labml.configs import option
20from labml_helpers.module import Module
21from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
22from labml_nn.normalization.deep_norm import DeepNormTransformerLayer
23from labml_nn.transformers import MultiHeadAttention
24from labml_nn.transformers.feed_forward import FeedForward

#

自己回帰モデル

これはDeepNormを使用する自己回帰変換モデルです。

27class AutoregressiveTransformer(Module):

#

n_tokens ボキャブラリ内のトークンの数です
d_model は埋め込みサイズ
n_layers 変圧器層の数です
layer レイヤーです。n_layers トランスフォーマーにはこれのコピーを使います

。

34    def __init__(self, n_tokens: int, d_model: int, n_layers: int, layer: DeepNormTransformerLayer):

#

41        super().__init__()

#

n_layers 層付き変圧器

43        self.transformer = nn.Sequential(*[copy.deepcopy(layer) for _ in range(n_layers)])

#

トークン埋め込みレイヤー

46        self.emb = nn.Embedding(n_tokens, d_model)

#

読み出し層

48        self.readout = nn.Linear(d_model, n_tokens)

#

x 形状の入力トークンです [seq_len, batch_size]

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

#

トークンの埋め込みを入手

55        x = self.emb(x)

#

トランスエンコーダー

57        x = self.transformer(x)

#

ロジットを取得

59        x = self.readout(x)

#

結果を返す

62        return x, None

#

コンフィギュレーション

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

65class Configs(NLPAutoRegressionConfigs):

#

モデル

74    model: AutoregressiveTransformer

#

レイヤー数

77    n_layers: int = 32

#

$α$ $β$ そしてディープノーム用

80    deep_norm_alpha: float
81    deep_norm_beta: float

#

注目されているヘッドの数

84    n_heads: int = 4

#

埋め込みサイズ

86    d_model: int = 64

#

各アテンションヘッドのサイズ

88    d_k: int = 16

#

計算 $α$

$α = (2 M)^{\frac{1}{4}}$

91@option(Configs.deep_norm_alpha)
92def _deep_norm_alpha(c: Configs):

#

98    return (2. * c.n_layers) ** (1. / 4.)

#

計算 $β$

$β = (8 M)^{- \frac{1}{4}}$

101@option(Configs.deep_norm_beta)
102def _deep_norm_beta(c: Configs):

#

108    return (8. * c.n_layers) ** -(1. / 4.)

#

モデルを初期化

111@option(Configs.model)
112def _model(c: Configs):

#

116    m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.n_layers,
117                                  DeepNormTransformerLayer(d_model=c.d_model,
118                                                           deep_norm_alpha=c.deep_norm_alpha,
119                                                           deep_norm_beta=c.deep_norm_beta,
120                                                           feed_forward=FeedForward(d_model=c.d_model,
121                                                                                    d_ff=c.d_model * 4),
122                                                           self_attn=MultiHeadAttention(c.n_heads, c.d_model,
123                                                                                        dropout_prob=0.0)))
124
125    return m.to(c.device)

#

実験を作成して実行する

128def main():

#

実験を作成

133    experiment.create(name="deep_norm", writers={'screen', 'web_api'})

#

コンフィグの作成

135    conf = Configs()

#

オーバーライド設定

137    experiment.configs(conf, {

#

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

139        'tokenizer': 'character',

#

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

141        'prompt_separator': '',

#

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

143        'prompt': 'It is ',

#

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

145        'text': 'tiny_shakespeare',

#

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

148        'seq_len': 256,

#

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

150        'epochs': 32,

#

バッチサイズ $16$

152        'batch_size': 16,

#

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

154        'inner_iterations': 10,

#

レイヤー数

157        'n_layers': 50,

#

ウォームアップなしの Adam オプティマイザー

161        'optimizer.optimizer': 'Adam',
162        'optimizer.learning_rate': 1.25e-4,
163    })

#

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

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

#

実験を始める

169    with experiment.start():

#

トレーニングを実行

171        conf.run()

#

175if __name__ == '__main__':
176    main()