From 5f07d2bf91d9aa43824209a160bd3ac69f43208f Mon Sep 17 00:00:00 2001 From: Anchor Yu <91590308+1azyking@users.noreply.github.com> Date: Wed, 25 Dec 2024 16:37:56 +0800 Subject: [PATCH] Update train-energy-hessian.md Signed-off-by: Anchor Yu <91590308+1azyking@users.noreply.github.com> --- doc/model/train-energy-hessian.md | 122 ++++++++++++++++++++++++++++-- 1 file changed, 117 insertions(+), 5 deletions(-) diff --git a/doc/model/train-energy-hessian.md b/doc/model/train-energy-hessian.md index 02b3c52052..fb4b72e9c3 100644 --- a/doc/model/train-energy-hessian.md +++ b/doc/model/train-energy-hessian.md @@ -4,10 +4,6 @@ **Supported backends**: PyTorch {{ pytorch_icon }} ::: -:::{warning} -The model trained with Hessian cannot be frozen. If freezing is forced, the model will be treated as a standard energy model, and the frozen model will no longer be able to output Hessian predictions. -::: - To train a model that takes Hessian matrices, i.e., the second order derivatives of energies w.r.t coordinates as input, you only need to prepare full Hessian matrices and modify the `loss` section to define the Hessian-specific settings, keeping other sections the same as the normal energy model's input script. ## Energy Hessian Loss @@ -32,7 +28,7 @@ The options {ref}`start_pref_e `, {ref}`limit_pref If one does not want to train with virial, then he/she may set the virial prefactors {ref}`start_pref_v ` and {ref}`limit_pref_v ` to 0. -## Hessian format in PyTorch +## Hessian Format in PyTorch In the PyTorch backend, Hessian matrices are listed in `hessian.npy` files, and the data format may contain the following files: @@ -52,3 +48,119 @@ This system contains `Nframes` frames with the same atom number `Natoms`, the to | hessian | Hessian matrices | hessian.npy | eV/Å^2 | Nframes \* (Natoms \* 3 \* Natoms \* 3) | Second-order derivatives of energies w.r.t coordinates. | Note that the `hessian.npy` should contain the **full** Hessian matrices with shape of `(3Natoms * 3Natoms)` for each frame, rather than the upper or lower triangular matrices with shape of `(3Natoms * (3Natoms + 1) / 2)` for each frame. + +## Train the Model + +There are two approaches to training a Hessian model. The first method involves training the model from scratch using the same command as in the `ener` mode within the PyTorch backend: + +::::{tab-set} + +:::{tab-item} PyTorch {{ pytorch_icon }} + +```bash +dp --pt train input.json +``` + +::: + +:::: + +The second approach is to train a Hessian model from a pretrained energy model, following the same command as the `finetune` strategy within the PyTorch backend: + +::::{tab-set} + +:::{tab-item} PyTorch {{ pytorch_icon }} + +```bash +dp --pt train input.json --finetune pretrained_energy.pt +``` + +::: + +:::: + +The detailed loss can be found in `lcurve.out`: + +``` +# step rmse_val rmse_trn rmse_e_val rmse_e_trn rmse_f_val rmse_f_trn rmse_h_val rmse_h_trn lr + 0 1.05e+02 2.28e+01 2.11e-01 1.59e+00 3.25e+00 3.37e-01 6.00e+00 6.37e+00 1.0e-03 + 200 1.86e+01 3.23e+01 9.24e-03 1.54e-01 2.51e-01 4.70e-01 5.31e+00 9.05e+00 1.0e-03 + 400 2.69e+01 2.98e+01 1.03e-01 1.07e-01 5.67e-01 4.17e-01 6.35e+00 8.47e+00 1.0e-03 + 600 2.00e+01 1.90e+01 7.23e-02 6.90e-03 3.35e-01 2.58e-01 5.37e+00 5.41e+00 1.0e-03 + 800 1.68e+01 1.48e+01 4.06e-02 2.27e-01 2.35e-01 1.98e-01 4.76e+00 4.24e+00 1.0e-03 + 1000 1.70e+01 1.81e+01 3.90e-01 1.66e-01 2.02e-01 1.99e-01 4.98e+00 5.37e+00 1.0e-03 +``` + +## Test the Model + +:::{warning} +A model trained with Hessian cannot be frozen. If freezing is enforced, the model will be treated as a standard energy model, and the frozen one will no longer be able to output Hessian predictions. +::: + +If one do freeze and test a Hessian model using the commands: + +::::{tab-set} + +:::{tab-item} PyTorch {{ pytorch_icon }} + +```bash + +dp --pt freeze -o frozen_model.pth + +dp --pt test -m frozen_model.pth -s test_system -d ${output_prefix} -a -n 1 +``` + +::: + +:::: + +If `dp test -d ${output_prefix} -a` is specified, the output files will be the same as those in the `ener` mode, i.e., + +``` +${output_prefix}.e.out ${output_prefix}.e_peratom.out ${output_prefix}.f.out +${output_prefix}.v.out ${output_prefix}.v_peratom.out +``` + +If one intends to use the trained model for Hessian predictions, then he/she is supposed to test the model directly without performing a freezing operation: + +::::{tab-set} + +:::{tab-item} PyTorch {{ pytorch_icon }} + +```bash + +dp --pt test -m model.pt -s test_system -d ${output_prefix} -a -n 1 +``` + +::: + +:::: + +If `dp test -d ${output_prefix} -a` is specified, the predicted Hessian for each frame are output in an additional file in the working directory: + +``` +${output_prefix}.h.out +``` + +For `*.h.out.*`, it contains matrix with shape of `(2, n_hess)`: + +``` +# frame - 0: data_h pred_h (3Na*3Na matrix in row-major order) +5.897392891323943331e+01 2.909700516268236825e+01 +-7.682282297964052376e+00 2.535680817045881774e+00 +-1.266442953072092514e+01 -2.127310638041492652e+01 +5.442541716174009031e-02 7.202825779190234756e-02 +5.198263170894957939e-05 -8.110080221576332349e-02 +7.443552765043950914e-02 -2.248597801730128215e-02 +1.029910175689553675e+00 1.938646932394622047e-03 +1.213862217511276764e+00 5.344132558814301825e-02 +-1.221943904909605250e+00 1.602557574981743893e-01 +``` + +The full Hessian matrices are stored in a flattened form in the row-major order. Here, `n_hess` is the total number of Hessian matrix elements across all frames, calculated as: + +```math +n_\text{hess} = \sum_{i} 3N_{\text{atom}, i}*3N_{\text{atom}, i} +``` + +where $N_{\text{atom}, i}$ represents the number of atoms in the $i^{\text{th}}$ frame.