付録1: 訓練用制御ファイルの高度な使用法

訓練用制御ファイルの全パラメータ#

{
    "common_config": {
        "total_epoch": null,
        "batch_size": 128,
        "reload_model": null,
        "lr_scheduling": "linear_warmup",
        "max_energy": 1e38,
        "max_forces": 1e38
    },
    "mtp_config": {
        "model_version": "v1.0",
        "pretrained_model": "ALL_ELEMENTS_SMALL_6"
    },
    "multistage_config": {
        "lr_list": [0.1, 0.01],
        "epoch_list": [80, 160],
        "energy_loss_coeff_list": [1e-05, 1.0],
        "forces_loss_coeff_list": [10.0, 0.1],
        "stress_loss_coeff_list": [1e-05, 10.0],
        "loss_func_list": ["mse", "mse"],
        "use_calib_stage": true,
        "calibration_lr": 0.01,
        "calibration_epoch": 80
    }
}

例の説明#

上記の制御ファイルの例を用いて、訓練ジョブがどのように動作するかを説明します。

最初に、”mtp_config”に従ってlight PFPモデルが初期化されます。コードはMTPアーキテクチャであることを識別し、指定された事前学習済みモデル”ALL_ELEMENTS_SMALL_6”でMTPモデルが初期化されます。

Note

“ORGANIC”など特定の事前学習済みモデルは、元素数に制限があります。サポートされていない元素でモデルを訓練しようとすると、エラーが発生します。
複数段階の訓練が開始されます。最初の段階では、学習率は0.1で、エネルギー/力/応力損失の係数はそれぞれ1e-5、10.0、1e-5です。主に力に対する性能を向上させるため、エネルギーと応力の係数はほぼ0に設定されています。このような戦略によって、高品質のlight PFPモデルを効果的に得ることができます。データセットは最初にスキャンされ、最大エネルギー/力を超えるデータポイントが削除され、次にトレーニングセットとテストセットに分割されます。トレーニングセットは、128個の構造を含むミニバッチでMTPモデルに渡されます。
最初の段階が終了した後、異なる損失係数、学習率、学習率スケジューラで第2段階が開始されます。最初の段階の学習でパラメータ空間の適当な領域に収束しているため、この段階では学習率は最初の段階よりも小さい0.01に設定されています。
“use_calib_stage”がtrueであるため、最後にキャリブレーションステージが実行されます。第2段階のトレーニングログが分析され、エネルギー/力/応力の損失の寄与がほぼ等しくなるように、3つの損失係数が自動的に割り当てられます。

パラメータの説明#

共通設定

訓練タスクの基本的な引数です。

total_epoch (int, Optional):

複数段階の訓練プロセス全体が実行されるエポック数。stage1/stage2/calibrationステージのそれぞれに対し、0.25:0.5:0.25の比率で分割されます。
batch_size (int, Optional):

1回のパラメータ更新に使用されるトレーニング構造の数。
reload_model (str, Optional):

モデルIDが指定されている場合、指定されたモデルから訓練ジョブが開始されます。空の場合、モデルのパラメータがランダムに初期化されます。
lr_scheduling (str, Optional):

学習率スケジューラ。”step”、”linear_warmup”、”None”をサポートしています。
scheduler_step_size` (int, Optional):

スケジューラのステップサイズ。デフォルト値はint(epoch / 25)です。
max_energy (float, Optional):

訓練データセットで許容される1原子あたりの最大エネルギー。エネルギーがこの値より大きい場合、構造はスキップされます。
max_forces (float, Optional):

訓練データセットで許容される最大の力。1原子に作用する力がこの値より大きい場合、構造はスキップされます。

モデル: MTP設定

MPTモデルの複雑性を定義するための引数です。

model_version (str, Optional):

MTPモデルのバージョン。現在、”v1.0”のみがサポートされています。
pretrained_model (str, Optional):

事前学習済みモデルを使用します。サポートされている事前学習済みモデルは、 (a) 事前学習済みモデルセクションにリストされています。空の場合、事前学習済みモデルは使用されません。デフォルトは “” です。

Note

このフィールドが指定されている場合、mtp_configセクションのmodel_version以外の残りのフィールドは無視されます。

[オプション] 次の引数はスクラッチでMTPモデルを訓練する場合にのみ使用されます（詳細については、 6.2 スクラッチでの訓練セクションを参照してください）。

rc (float, Optional):

1原子の局所環境を定義するためのカットオフ距離。
levmax (int, Optional):

モーメント記述子の複雑性を制御する引数。 “levmax”パラメータが大きいほど、MTPモデルはより複雑になります。
n_q (int, Optional):

MTPの動径関数を記述するために使用されるChebyshev多項式の数。
moment_mu_cost (int, Optional):

モーメントM(𝜇, 𝜈)のレベルを計算する際の𝜇のコスト係数。コストはモーメントのmuとnuの値には依存しません。
moment_init_cost (int, Optional):

モーメントM(𝜇, 𝜈)のレベルを計算する際に使用される初期コスト。
moment_nu_cost (int, Optional):

モーメントM(𝜇, 𝜈)のレベルを計算する際の𝜈のコスト係数。
nn_dim_list (List[int], Optional):

ニューラルネットワークの各層のニューロン数。リストの長さは層の数を表します。
デフォルトは []で、この場合は線形関数が readout 関数に使用されます。
空でない場合、ニューラルネットワークが readout 関数に使用されます。

Note

推奨される隠れ層の数は2〜3層で、各層のニューロン数は16〜128です。

複数段階の訓練

訓練タスクの引数です。 “lr_list”、”epoch_list”、”energy_loss_coeff_list”、”forces_loss_coeff_list”、”stress_loss_coeff_list”、”loss_func_list”の要素数を等しくする必要があります。

lr_list (List[float], Optional):

学習率のリスト。i番目の要素はi番目のステージの学習率です。
epoch_list (List[int], Optional):

各ステージのトレーニングエポック数。
energy_loss_coeff_list (List[float], Optional):

各ステージのエネルギー損失の係数。
forces_loss_coeff_list (List[float], Optional):

各ステージの力損失の係数。
stress_loss_coeff_list (List[float], Optional):

各ステージの応力損失の係数。
loss_func_list (List[str], Optional):

各ステージでエネルギー/力/応力損失を評価するための損失関数。”mse”、”l1”をサポートします。
use_calib_stage (bool, Optional):

キャリブレーションステージを使用するかどうか。
calibration_lr (float, Optional):

キャリブレーションステージの学習率。
calibration_epoch (int, Optional):

キャリブレーションステージのエポック数。