deepmodelingcommunity/DPA-2.4-7M

Description

DPA-2.4-7M is trained using a multitask strategy on the OpenLAM datasets. The benchmark results can be found at the LAMBench website.

Please ensure the usage of the correct code version (v3.1.0a0) corresponding to this model.

How to use

Installation

To use the pretrained model, you need to first install the corresponding version of DeePMD-kit. You can try easy installation:

pip install torch torchvision torchaudio
pip install git+https://github.com/deepmodeling/deepmd-kit@v3.1.0a0

For other installation options, please visit the Releases page page to download the off-line package for deepmd-kit v3.1.0a0, and refer to the official documentation for off-line installation instructions.

Use the pretrained model

The pretrained model can be used directly for prediction tasks, such as serving as an ASE calculator.

This model is pretrained in a multi-task manner, featuring a unified backbone (referred to as the unified descriptor) and several fitting nets for different datasets. For detailed information, refer to the DPA-2 paper.

The first step involves selecting a specific fitting net from the model to make predictions. To list the available fitting nets (model-branch) in this pretrained model, use the following command:

dp --pt show dpa-2.4.pt model-branch

This will generate an output similar to the following:

Available model branches are ['Domains_Alloy', 'Domains_Anode', 'Domains_Cluster', 
'Domains_Drug', 'Domains_FerroEle', 'Domains_SSE_PBE', 'Domains_SemiCond', 'H2O_H2O_PD', 
'Metals_AlMgCu', 'Metals_Cu', 'Metals_Sn', 'Metals_Ti', 'Metals_V', 'Metals_W', 'Others_C12H26',
'Others_HfO2', 'Domains_SSE_PBESol', 'Domains_Transition1x', 'H2O_H2O_DPLR', 
'H2O_H2O_PBE0TS_MD', 'H2O_H2O_PBE0TS', 'H2O_H2O_SCAN0', 'Metals_AgAu_PBED3', 'Others_In2Se3', 
'MP_traj_v024_alldata_mixu', 'Alloy_tongqi', 'SSE_ABACUS', 'Hybrid_Perovskite', 
'solvated_protein_fragments', 'Electrolyte', 'ODAC23', 'Alex2D', 'Omat24', 'SPICE2', 'OC20M', 
'OC22', 'Organic_Reactions', 'RANDOM'], where 'RANDOM' means using a randomly initialized 
fitting net.

Select a fitting net that closely matches your system. Ensure that the elements in your system are included in the corresponding fitting net if you are conducting direct predictions or simulations. For more information on the pretrained datasets, refer to below. Note: model branches including "Metals_Cu","Others_C12H26","H2O_H2O_DPLR","H2O_H2O_PBE0TS_MD","H2O_H2O_PBE0TS", and "H2O_H2O_SCAN0" are deprecated (training datasets were removed during the training process).

Assuming you choose H2O_H2O-PD, you can first freeze the model branch from the multi-task pretrained model:

dp --pt freeze -c dpa-2.4-7M.pt -o frozen_model.pth --head H2O_H2O-PD

Then you can use the following Python code for prediction or optimization:

## Compute potential energy
from ase import Atoms
from deepmd.calculator import DP as DPCalculator
dp = DPCalculator("frozen_model.pth")
water = Atoms('H2O', positions=[(0.7601, 1.9270, 1), (1.9575, 1, 1), (1., 1., 1.)], cell=[100, 100, 100])
water.calc = dp
print(water.get_potential_energy())
print(water.get_forces())

## Run BFGS structure optimization
from ase.optimize import BFGS
dyn = BFGS(water)
dyn.run(fmax=1e-6)
print(water.get_positions())

Zero-Shot Inference

Given that energy labels calculated via DFT can vary by an arbitrary constant, LAMs are consistently used to predict the energy difference between the label and a dummy model that estimates potential energy solely based on the chemical formula. Specifically, the dummy model is defined as

E = Σₑ nₑ · b̂ₑ;

where E is the predicted energy, nₑ is the number of atoms of element e, and b̂ₑ is the per-atom energy of element e obtained via least-squares fitting:

b̂ = argmin Σᵢ ( yᵢ − Σₑ nᵢ,ₑ bₑ )²,

with nᵢ,ₑ denoting the number of atoms of element e in sample i, and yᵢ the corresponding DFT energy label. When applying the pretrained model to a new downstream system, it is advisable to adjust the energy bias term to better align with the DFT energy labels of the downstream system. This adjustment can enhance prediction accuracy without necessitating further training of the neural network parameters.

To be specific, given a new downstream system, we use dp --pt change-bias to do zero-shot procedure, which will inherit the neural network parameters of the pretrained multitask model, but updating the energy bias to better align with the downstream system.

dp --pt change-bias dpa-2.4-7M.pt -s <your_system> --model-branch Omat24

Fine-tuning

To fine-tune a pretrained model, you need to ensure that the model configuration (model/descriptor and model/fitting_net) matches exactly with the pretrained model. The input_finetune.json file provided on this page contains the same model configuration as the pretrained model.

To prepare for fine-tuning, replace the model/type_map and training/training(validation)_data sections in input_finetune.json with your specific dataset information, just the same as from scratch. Additionally, adjust the learning rate for fine-tuning, typically a small learning rate such as 1e-4 for learning_rate/start_lr is recommended. One can inherit a pretrained fitting net by using the --model-branch argument.

dp --pt train input_finetune.json --finetune dpa-2.4-7M.pt --model-branch Omat24