DeePMD-kit supports seven backends: TensorFlow, TensorFlow 2, PyTorch-TorchScript, PyTorch-Exportable, JAX, Paddle, and the NumPy-based DP reference backend. To use DeePMD-kit, you must install at least one backend. Each backend does not support all features. In the documentation, TensorFlow and TensorFlow 2 share {{ tensorflow_icon }}, while PyTorch-TorchScript and PyTorch-Exportable share {{ pytorch_icon }}. JAX {{ jax_icon }}, Paddle {{ paddle_icon }}, and DP {{ dpmodel_icon }} use separate icons. Support notes spell out the exact backend variant when the two implementations in a framework family differ.
- Model filename extension:
.pb - Checkpoint filename extension:
.meta,.index,.data-00000-of-00001
TensorFlow 2.8 is the first version to support Python 3.10.
DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (tf.compat.v1) in the graph mode.
- Model filename extension:
.savedmodeltf - Checkpoint directory extension:
.tf2
The TensorFlow 2 backend uses the TensorFlow v2 eager API. Select it with
dp --tf2 (alias dp --tensorflow2). It supports training, including
multi-task training and fine-tuning. Freezing, compression, and testing use a
.savedmodeltf export and therefore require graph-traceable model code.
Training stores checkpoints in a directory named after the save_ckpt prefix
with .tf2 appended, such as model.ckpt.tf2.
For training, set
{ref}training.enable_compile <training/enable_compile> to true to enable
XLA compilation of the formatted lower-forward path. Setting
DP_JIT enables the same model-level default and also
applies it to SavedModel export. Depending on the workload, compilation may
improve or reduce performance.
- Model filename extension:
.pth - Checkpoint filename extension:
.pt
PyTorch 2.1 or above is required. Select this backend
with dp --pt. It uses TorchScript for most frozen models; DPA4/SeZM uses a
separate AOTInductor export path. Because PyTorch has deprecated TorchScript,
DeePMD-kit will deprecate this backend and replace it with PyTorch-Exportable.
While .pth and .pt are the same in the PyTorch package, they have different
meanings in DeePMD-kit: .pth stores a frozen model, while .pt stores a
training checkpoint.
- Model filename extensions:
.pte,.pt2 - Checkpoint filename extension:
.pt
Select this backend with dp --pt-expt (alias
dp --pytorch-exportable). It uses PyTorch with the backend-independent model
implementation and supports training, including multi-task training and
fine-tuning, freezing, change-bias, and testing. Compression support and export
requirements are documented on the corresponding descriptor pages. Training
can read LMDB datasets, and Python inference can use the optional vesin
neighbor-list implementation.
Freezing exports a torch.export model. The dense neighbor-list lower form
normally uses .pte, while the graph lower form uses an AOTInductor .pt2
package. Use --lower-kind graph to request graph-native export for an eligible
model; graph-capable DPA models may select that form automatically. The .pt
checkpoint format uses DP-model parameter names ending in .w and .b, which
allows DeePMD-kit to distinguish it from a PyTorch-TorchScript checkpoint,
whose parameter names end in .matrix and .bias.
The .pt2 suffix identifies an AOTInductor package, but not its lower-input
ABI. A DPA4/SeZM model frozen with dp --pt freeze normally uses the legacy
edge_vec ABI (lower_input_kind: edge_vec); its deepspin virtual-atom variant
uses the dense nlist ABI instead. A graph model frozen with
dp --pt-expt freeze --lower-kind graph uses the NeighborGraph ABI
(lower_input_kind: graph). All variants are loaded for inference by the
PyTorch-Exportable runtime, which reads this metadata to select the correct
input path. The --lower-kind option controls only the PyTorch-Exportable
freeze route; see the DPA4 export documentation
for the separate PyTorch-TorchScript route.
- DeepEval model filename extensions:
.hlo,.savedmodel - Checkpoint and lossless serialization extension:
.jax
JAX 0.4.33 or above is required.
Both .hlo and .jax are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions.
.savedmodel is the TensorFlow SavedModel format generated by JAX2TF, which needs the installation of TensorFlow.
Only the .savedmodel format supports C++ inference, which needs the TensorFlow C++ interface.
The model is device-specific, so that the model generated on the GPU device cannot be run on the CPUs.
JAX supports training with dp --jax train; training checkpoints use the
.jax extension. Freezing can write a DeepEval-compatible .hlo or
.savedmodel model, or a lossless .jax serialization for checkpoint
round-tripping and JAX-MD. The normal dp test/DeepPot route does not load
.jax serializations.
- Model filename extensions:
.jsonand.pdiparams - Checkpoint filename extension:
.pd
Paddle version 3.0 or above is required.
The .pd extension is used for model checkpoint storage, which is commonly utilized during training and testing in Python. The .json extension is for the model's computational graph in PIR representation, while the .pdiparams extension stores model parameters. Both .json and .pdiparams files are exported together and used in model freezing and Python/C++ inference.
Note
This backend is only for development and should not take into production.
- Model filename extension:
.dp,.yaml,.yml
DP is a reference backend for development, which uses pure NumPy to implement models without using any heavy deep-learning frameworks.
Due to the limitation of NumPy, it doesn't support gradient calculation and thus cannot be used for training.
As a reference backend, it is not aimed at the best performance, but only the correct results.
The DP backend has two formats, both of which are backend-independent:
The .dp format uses HDF5 to store model serialization data, which has good performance.
The .yaml or .yml use YAML to save the data as plain texts, which is easy to read for human beings.
Only Python inference interface can load these formats.
NumPy 1.21 or above is required.
When training and freezing a model, use dp --tf, dp --tf2, dp --pt,
dp --pt-expt, dp --jax, or dp --pd in the command line to switch the
backend.
When doing inference, DeePMD-kit detects the backend from the model filename.
For example, when the model filename ends with .pb (the ProtoBuf file), DeePMD-kit will consider it using the TensorFlow backend.
The same detection covers TensorFlow 2 .savedmodeltf models and
PyTorch-Exportable .pte and .pt2 runtime formats. In particular, .pt2
selects the PyTorch-Exportable inference loader even when the file was produced
by the DPA4/SeZM dp --pt freeze route described above; the archive metadata
then selects its edge_vec, dense nlist, or NeighborGraph ABI.
If a model is supported by two backends, one can use dp convert-backend to convert the model file between these two backends.