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
19class Configs(MLMConfigs):
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
論文はアクティベーションを示唆しています
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
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,
シーケンスの長さは 短いシーケンス長を使用してトレーニングを高速化します。そうしないと、MLMモデルのトレーニングには永遠に時間がかかります
。81 'seq_len': 32,
1024 エポックのトレーニングを行います。
84 'epochs': 1024,
エポックごとにトレーニングと検証を切り替える
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,
保存および読み込み用のモデルを設定する
105 experiment.add_pytorch_models({'model': conf.model})
実験を始める
108 with experiment.start():
トレーニングを実行
110 conf.run()
114if __name__ == '__main__':
115 main()