#

ロータリー・ポジショナル・エンベディング (RoPE) 実験

これは、回転式位置埋め込み（RoPE）を使用してトランスフォーマーモデルをトレーニングするための注釈付きPyTorch実験です。

12from labml import experiment
13from labml.configs import calculate
14from labml_nn.transformers import TransformerConfigs
15from labml_nn.transformers.rope.experiment import Configs as RoPEConfigs

#

ロータリーPE注意

#

20class Configs(RoPEConfigs):  # , ArithmeticAutoregression):
21    pass

#

24def _rotary_value_pe_mha(c: TransformerConfigs):
25    from labml_nn.transformers.rope.value_pe import RotaryValuePEMultiHeadAttention
26    return RotaryValuePEMultiHeadAttention(c.n_heads, c.d_model, 1., 1.)

#

設定オプション

30calculate(TransformerConfigs.encoder_attn, 'rotary_value', _rotary_value_pe_mha)
31calculate(TransformerConfigs.decoder_attn, 'rotary_value', _rotary_value_pe_mha)
32calculate(TransformerConfigs.decoder_mem_attn, 'rotary_value', _rotary_value_pe_mha)

#

35def main():

#

実験を作成

37    experiment.create(name="rotary_shakespeare", comment="rotary value", writers={'screen', 'labml'})

#

コンフィグの作成

39    conf = Configs()

#

オーバーライド設定

41    experiment.configs(conf, {

#

固定位置埋め込みなし

43        'transformer.src_embed': 'no_pos',
44        'transformer.tgt_embed': 'no_pos',

#

RoPE 付きエンコーダ

47        'transformer.encoder_attn': 'rotary_value',

#

'トランスフォーマー.encoder_attn ':' ロータリー '、

#

51        'model': 'rotary_pe_transformer',

#

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

54        'tokenizer': 'character',

#

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

56        'prompt_separator': '',

#

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

58        'prompt': 'It is ',

#

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

60        'text': 'tiny_shakespeare',

#

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

63        'seq_len': 512,

#

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

65        'epochs': 24,

#

バッチサイズ $4$

67        'batch_size': 16,

#

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

70        'inner_iterations': 4,

#

モデルサイズ

73        'd_model': 128,
74        'transformer.ffn.d_ff': 512,
75        'transformer.n_heads': 4,
76        'transformer.dropout': 0.0,

#

Adam オプティマイザーを使う

79        'optimizer.optimizer': 'Adam',
80        'optimizer.learning_rate': 2.5e-4,
81
82        'dataloader_shuffle_with_replacement': True
83    })

#

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

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

#

実験を始める

89    with experiment.start():

#

トレーニングを実行

91        conf.run()

#

95if __name__ == '__main__':
96    main()