#

MLP ミキサー実験

これは MLP ミキサーモデルをトレーニングするための注釈付き PyTorch 実験です。

12from labml import experiment
13from labml.configs import option
14from labml_nn.transformers import TransformerConfigs
15from labml_nn.transformers.configs import FeedForwardConfigs
16from labml_nn.transformers.mlm.experiment import TransformerMLM, Configs as MLMConfigs

#

コンフィギュレーション

これは、MLMConfigs マスク言語モデルの実験を定義するところから受け継がれています。

19class Configs(MLMConfigs):

#

MLP 用の設定可能なフィードフォワードネットワーク

29    mix_mlp: FeedForwardConfigs

#

MLP 構成のミキシング

32@option(Configs.mix_mlp)
33def _mix_mlp_configs(c: Configs):

#

38    conf = FeedForwardConfigs()

#

MLPのサイズはトークン全体に適用されるため、シーケンスの長さになります

40    conf.d_model = c.seq_len

#

$GE LU$ 論文はアクティベーションを示唆しています

42    conf.activation = 'GELU'

#

45    return conf

#

変圧器構成

48@option(Configs.transformer)
49def _transformer_configs(c: Configs):

#

設定可能なトランス実装を使用しています

56    conf = TransformerConfigs()

#

埋め込みやロジットの生成に使用するボキャブラリーサイズを設定

58    conf.n_src_vocab = c.n_tokens
59    conf.n_tgt_vocab = c.n_tokens

#

埋め込みサイズ

61    conf.d_model = c.d_model

#

アテンションモジュールをMLPMixerに変更

63    from labml_nn.transformers.mlp_mixer import MLPMixer
64    conf.encoder_attn = MLPMixer(c.mix_mlp.ffn)

#

67    return conf

#

70def main():

#

実験を作成

72    experiment.create(name="mlp_mixer_mlm")

#

コンフィグの作成

74    conf = Configs()

#

オーバーライド設定

76    experiment.configs(conf, {

#

バッチサイズ

78        'batch_size': 64,

#

シーケンスの長さは $32$ 短いシーケンス長を使用してトレーニングを高速化します。そうしないと、MLMモデルのトレーニングには永遠に時間がかかります

。

81        'seq_len': 32,

#

1024 エポックのトレーニングを行います。

84        'epochs': 1024,

#

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

87        'inner_iterations': 1,

#

変圧器構成

90        'd_model': 128,
91        'transformer.ffn.d_ff': 256,
92        'transformer.n_heads': 8,
93        'transformer.n_layers': 6,
94        'transformer.ffn.activation': 'GELU',

#

ミキサー MLP 隠れレイヤーサイズ

97        'mix_mlp.d_ff': 128,

#

Noam オプティマイザを使う

100        'optimizer.optimizer': 'Noam',
101        'optimizer.learning_rate': 1.,
102    })

#

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

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

#

実験を始める

108    with experiment.start():

#

トレーニングを実行

110        conf.run()

#

114if __name__ == '__main__':
115    main()