#

階層型トランスフォーマーはより効率的な言語モデルの実験

これは、単純なトランスフォーマー自己回帰NLPタスクのトレーニングループと構成に基づいています。

14import math
15from typing import List
16
17import torch
18from torch import nn
19
20from labml import experiment
21from labml.configs import option
22from labml_helpers.module import Module
23from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
24from labml_nn.transformers.hour_glass import HourGlass
25from labml_nn.transformers.positional_encoding import PositionalEncoding

#

自己回帰言語モデル

28class AutoregressiveTransformer(Module):

#

n_tokens はボキャブラリーサイズです
d_model トークンの埋め込みのサイズです
dropout は脱落確率です
hour_glass 砂時計モデルです

33    def __init__(self, n_tokens: int, d_model: int, dropout: float, hour_glass: HourGlass):

#

40        super().__init__()

#

トークンの埋め込み

42        self.embedding = nn.Embedding(n_tokens, d_model)

#

位置埋め込みを修正しました。

📝 公式論文の実装には比較的注意が必要

48        self.pos_embedding = PositionalEncoding(d_model, dropout)

#

砂時計モデル

50        self.hour_glass = hour_glass

#

最終的な埋め込みを正規化するには

52        self.norm = nn.LayerNorm([d_model])

#

埋め込みサイズ

54        self.d_model = d_model

#

ロジットを予測するための最後の線形レイヤー

56        self.output = nn.Linear(d_model, n_tokens)

#

x 形状のトークンインデックスが付いたテンソルです [seq_len, batch_size]

58    def __call__(self, x: torch.Tensor):

#

埋め込みを入手

63        x = self.embedding(x)

#

位置埋め込みを追加

66        if self.pos_embedding is not None:
67            x = self.pos_embedding(x * math.sqrt(self.d_model))

#

砂時計

70        x = self.hour_glass(x)

#

ロジットを取得

73        output = self.output(self.norm(x))

#

ロジットを返す

76        return output, None

#

コンフィギュレーション

これは、単純なトランスフォーマー自己回帰NLPタスクのトレーニングループと構成から継承されています。

79class Configs(NLPAutoRegressionConfigs):

#

モデル

87    model: AutoregressiveTransformer

#

アテンションヘッドの数

89    n_heads: int = 8

#

脱落確率

91    dropout: float = 0.1

#

フィードフォワード隠れ層のサイズ

93    d_ff: int = 512

#

トークンの埋め込みサイズ

95    d_model: int = 256

#

短縮要因

97    shortening_factors: List[int] = [8, 4]

#

モデルを作成

100@option(Configs.model)
101def _model(c: Configs):

#

砂時計モデルを作成

107    hour_glass = HourGlass(c.n_heads, c.d_model, c.dropout, c.d_ff, c.shortening_factors)

#

自動回帰ラッパーの作成

109    m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.dropout, hour_glass).to(c.device)

#

112    return m

#

115def main():

#

実験を作成

117    experiment.create(name="hour_glass")

#

コンフィグの作成

119    conf = Configs()

#

オーバーライド設定

121    experiment.configs(conf, {

#

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

123        'tokenizer': 'character',

#

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

125        'prompt_separator': '',

#

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

127        'prompt': 'It is ',

#

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

129        'text': 'tiny_shakespeare',

#

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

132        'seq_len': 256,

#

$128$ 時代に合わせた列車

134        'epochs': 128,

#

バッチサイズ $32$

136        'batch_size': 32,

#

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

139        'inner_iterations': 10,

#

Noam オプティマイザを使う

142        'optimizer.optimizer': 'Noam',
143        'optimizer.learning_rate': 1.,

#

145    })

#

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

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

#

実験を始める

151    with experiment.start():

#

トレーニングを実行

153        conf.run()

#

157if __name__ == '__main__':
158    main()