#

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

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

12from labml import experiment
13from labml.configs import option, calculate
14from labml_nn.transformers import TransformerConfigs
15from labml_nn.transformers.basic.autoregressive_experiment import AutoregressiveTransformer, Configs

#

ロータリーPE注意

19def _rotary_pe_mha(c: TransformerConfigs):
20    from labml_nn.transformers.rope import RotaryPEMultiHeadAttention
21    return RotaryPEMultiHeadAttention(c.n_heads, c.d_model, 1.)

#

設定オプション

25calculate(TransformerConfigs.encoder_attn, 'rotary', _rotary_pe_mha)
26calculate(TransformerConfigs.decoder_attn, 'rotary', _rotary_pe_mha)
27calculate(TransformerConfigs.decoder_mem_attn, 'rotary', _rotary_pe_mha)

#

自己回帰モデルの作成と重みの初期化

30@option(Configs.model, 'rotary_pe_transformer')
31def _model(c: Configs):

#

35    m = AutoregressiveTransformer(c.transformer.encoder,
36                                  c.transformer.src_embed,
37                                  c.transformer.generator).to(c.device)
38
39    return m

#

42def main():

#

実験を作成

44    experiment.create(name="rotary_pe_transformer", writers={'screen'})

#

コンフィグの作成

46    conf = Configs()

#

オーバーライド設定

48    experiment.configs(conf, {

#

固定位置埋め込みなし

50        'transformer.src_embed': 'no_pos',
51        'transformer.tgt_embed': 'no_pos',

#

RoPE 付きエンコーダ

54        'transformer.encoder_attn': 'rotary',

#

57        'model': 'rotary_pe_transformer',

#

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

60        'tokenizer': 'character',

#

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

62        'prompt_separator': '',

#

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

64        'prompt': 'It is ',

#

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

66        'text': 'tiny_shakespeare',

#

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

69        'seq_len': 512,

#

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

71        'epochs': 32,

#

バッチサイズ $4$

73        'batch_size': 4,

#

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

76        'inner_iterations': 10,

#

モデルサイズ

79        'd_model': 128,
80        'transformer.ffn.d_ff': 512,
81        'transformer.n_heads': 16,
82        'transformer.dropout': 0.0,

#

Noam オプティマイザを使う

85        'optimizer.optimizer': 'Noam',
86        'optimizer.learning_rate': 1.,
87
88        'dataloader_shuffle_with_replacement': True
89    })

#

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

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

#

実験を始める

95    with experiment.start():

#

トレーニングを実行

97        conf.run()

#

101if __name__ == '__main__':
102    main()