You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
PyTorch-TorchScript, PyTorch-Exportable, JAX, Paddle, and the NumPy-based DP
7
+
reference backend.
6
8
To use DeePMD-kit, you must install at least one backend.
7
9
Each backend does not support all features.
8
-
In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, and Paddle {{ paddle_icon }} icons are used to mark whether a backend supports a feature.
10
+
In the documentation, TensorFlow and TensorFlow 2 share
11
+
{{ tensorflow_icon }}, while PyTorch-TorchScript and PyTorch-Exportable share
{{ dpmodel_icon }} use separate icons. Support notes spell out the exact backend
14
+
variant when the two implementations in a framework family differ.
9
15
10
16
### TensorFlow {{ tensorflow_icon }}
11
17
@@ -15,25 +21,91 @@ In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon
15
21
[TensorFlow](https://tensorflow.org) 2.8 is the first version to support Python 3.10.
16
22
DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`tf.compat.v1`) in the graph mode.
17
23
18
-
### PyTorch {{ pytorch_icon }}
24
+
### TensorFlow 2 {{ tensorflow_icon }}
25
+
26
+
- Model filename extension: `.savedmodeltf`
27
+
- Checkpoint directory extension: `.tf2`
28
+
29
+
The TensorFlow 2 backend uses the TensorFlow v2 eager API. Select it with
30
+
`dp --tf2` (alias `dp --tensorflow2`). It supports training, including
31
+
multi-task training and fine-tuning. Freezing, compression, and testing use a
32
+
`.savedmodeltf` export and therefore require graph-traceable model code.
33
+
Training stores checkpoints in a directory named after the `save_ckpt` prefix
34
+
with `.tf2` appended, such as `model.ckpt.tf2`.
35
+
36
+
For training, set
37
+
{ref}`training.enable_compile <training/enable_compile>` to `true` to enable
38
+
XLA compilation of the formatted lower-forward path. Setting
39
+
[`DP_JIT`](env.md#envvar-DP_JIT) enables the same model-level default and also
40
+
applies it to SavedModel export. Depending on the workload, compilation may
41
+
improve or reduce performance.
42
+
43
+
### PyTorch-TorchScript {{ pytorch_icon }}
19
44
20
45
- Model filename extension: `.pth`
21
46
- Checkpoint filename extension: `.pt`
22
47
23
-
[PyTorch](https://pytorch.org/) 2.1 or above is required.
24
-
While `.pth` and `.pt` are the same in the PyTorch package, they have different meanings in the DeePMD-kit to distinguish the model and the checkpoint.
48
+
[PyTorch](https://pytorch.org/) 2.1 or above is required. Select this backend
49
+
with `dp --pt`. It uses TorchScript for most frozen models; DPA4/SeZM uses a
50
+
separate AOTInductor export path. Because PyTorch has deprecated TorchScript,
51
+
DeePMD-kit will deprecate this backend and replace it with PyTorch-Exportable.
52
+
53
+
While `.pth` and `.pt` are the same in the PyTorch package, they have different
54
+
meanings in DeePMD-kit: `.pth` stores a frozen model, while `.pt` stores a
55
+
training checkpoint.
56
+
57
+
### PyTorch-Exportable {{ pytorch_icon }}
58
+
59
+
- Model filename extensions: `.pte`, `.pt2`
60
+
- Checkpoint filename extension: `.pt`
61
+
62
+
Select this backend with `dp --pt-expt` (alias
63
+
`dp --pytorch-exportable`). It uses PyTorch with the backend-independent model
64
+
implementation and supports training, including multi-task training and
65
+
fine-tuning, freezing, change-bias, and testing. Compression support and export
66
+
requirements are documented on the corresponding descriptor pages. Training
67
+
can read LMDB datasets, and Python inference can use the optional vesin
68
+
neighbor-list implementation.
69
+
70
+
Freezing exports a `torch.export` model. The dense neighbor-list lower form
71
+
normally uses `.pte`, while the graph lower form uses an AOTInductor `.pt2`
72
+
package. Use `--lower-kind graph` to request graph-native export for an eligible
73
+
model; graph-capable DPA models may select that form automatically. The `.pt`
74
+
checkpoint format uses DP-model parameter names ending in `.w` and `.b`, which
75
+
allows DeePMD-kit to distinguish it from a PyTorch-TorchScript checkpoint,
76
+
whose parameter names end in `.matrix` and `.bias`.
77
+
78
+
The `.pt2` suffix identifies an AOTInductor package, but not its lower-input
79
+
ABI. A DPA4/SeZM model frozen with `dp --pt freeze` normally uses the legacy
80
+
`edge_vec` ABI (`lower_input_kind: edge_vec`); its deepspin virtual-atom variant
81
+
uses the dense `nlist` ABI instead. A graph model frozen with
82
+
`dp --pt-expt freeze --lower-kind graph` uses the NeighborGraph ABI
83
+
(`lower_input_kind: graph`). The `dp --pt-expt compress` workflow can instead
84
+
export a compressed DPA-1 model through the compact canonical
85
+
`dpa1_canonical` ABI (`lower_input_kind: dpa1_canonical`) when the canonical
86
+
compression is eligible. All variants are loaded for inference by the
87
+
PyTorch-Exportable runtime, which reads this metadata to select the correct
88
+
input path. The `--lower-kind` option controls only the PyTorch-Exportable
89
+
freeze route; see the [DPA4 export documentation](model/dpa4.md#freeze-to-pt2)
90
+
for the separate DPA4/SeZM AOTInductor export route.
25
91
26
92
### JAX {{ jax_icon }}
27
93
28
-
-Model filename extension: `.xlo`, `.savedmodel`
29
-
- Checkpoint filename extension: `.jax`
94
+
-DeepEval model filename extensions: `.hlo`, `.savedmodel`
95
+
- Checkpoint and lossless serialization extension: `.jax`
30
96
31
97
[JAX](https://jax.readthedocs.io/) 0.4.33 or above is required.
32
-
Both `.xlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions.
98
+
Both `.hlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions.
33
99
`.savedmodel` is the TensorFlow [SavedModel format](https://www.tensorflow.org/guide/saved_model) generated by [JAX2TF](https://www.tensorflow.org/guide/jax2tf), which needs the installation of TensorFlow.
34
100
Only the `.savedmodel` format supports C++ inference, which needs the TensorFlow C++ interface.
35
101
The model is device-specific, so that the model generated on the GPU device cannot be run on the CPUs.
36
102
103
+
JAX supports training with `dp --jax train`; training checkpoints use the
104
+
`.jax` extension. Freezing can write a DeepEval-compatible `.hlo` or
105
+
`.savedmodel` model, or a lossless `.jax` serialization for checkpoint
106
+
round-tripping and JAX-MD. The normal `dp test`/`DeepPot` route does not load
107
+
`.jax` serializations.
108
+
37
109
### Paddle {{ paddle_icon }}
38
110
39
111
- Model filename extensions: `.json` and `.pdiparams`
@@ -64,12 +136,20 @@ NumPy 1.21 or above is required.
64
136
65
137
### Training
66
138
67
-
When training and freezing a model, you can use `dp --tf`, `dp --pt` or `dp --pd` in the command line to switch the backend.
139
+
When training and freezing a model, use `dp --tf`, `dp --tf2`, `dp --pt`,
140
+
`dp --pt-expt`, `dp --jax`, or `dp --pd` in the command line to switch the
141
+
backend.
68
142
69
143
### Inference
70
144
71
145
When doing inference, DeePMD-kit detects the backend from the model filename.
72
146
For example, when the model filename ends with `.pb` (the ProtoBuf file), DeePMD-kit will consider it using the TensorFlow backend.
147
+
The same detection covers TensorFlow 2 `.savedmodeltf` models and
148
+
PyTorch-Exportable `.pte` and `.pt2` runtime formats. In particular, `.pt2`
149
+
selects the PyTorch-Exportable inference loader even when the file was produced
150
+
by the DPA4/SeZM `dp --pt freeze` route described above; the archive metadata
151
+
then selects its `edge_vec`, dense `nlist`, NeighborGraph, or compact
> In the following context, we use the PyTorch backend as the example, while it also applies to other backends listed above.
6
+
> In the following context, we use the PyTorch-TorchScript backend as the example, while it also applies to other backends listed above.
7
7
8
8
If you'd like to create a new model that isn't covered by the existing DeePMD-kit library, but reuse DeePMD-kit's other efficient modules such as data processing, trainer, etc, you may want to read this section.
9
9
@@ -16,9 +16,9 @@ To incorporate your custom model you'll need to:
16
16
17
17
## Design a new component
18
18
19
-
With DeePMD-kit v3, we have expanded support to include two additional backends alongside TensorFlow: the PyTorch backend and the framework-independent backend (dpmodel). The PyTorch backend adopts a highly modularized design to provide flexibility and extensibility. It ensures a consistent experience for both training and inference, aligning with the TensorFlow backend.
19
+
With DeePMD-kit v3, we have expanded support to include two additional backends alongside TensorFlow: the PyTorch-TorchScript backend and the framework-independent backend (dpmodel). The PyTorch-TorchScript backend adopts a highly modularized design to provide flexibility and extensibility. It ensures a consistent experience for both training and inference, aligning with the TensorFlow backend.
20
20
21
-
The framework-independent backend is implemented in pure NumPy, serving as a reference backend to ensure consistency in tests. Its design pattern closely parallels that of the PyTorch backend.
21
+
The framework-independent backend is implemented in pure NumPy, serving as a reference backend to ensure consistency in tests. Its design pattern closely parallels that of the PyTorch-TorchScript backend.
22
22
23
23
### New descriptors
24
24
@@ -118,7 +118,7 @@ class SomeFittingNet(GeneralFitting):
118
118
119
119
### New models
120
120
121
-
The PyTorch backend's model architecture is meticulously structured with multiple layers of abstraction, ensuring a high degree of flexibility. Typically, the process commences with an atomic model responsible for atom-wise property calculations. This atomic model inherits from both the {py:class}`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` class and the {py:class}`torch.nn.Module` class.
121
+
The PyTorch-TorchScript backend's model architecture is meticulously structured with multiple layers of abstraction, ensuring a high degree of flexibility. Typically, the process commences with an atomic model responsible for atom-wise property calculations. This atomic model inherits from both the {py:class}`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` class and the {py:class}`torch.nn.Module` class.
122
122
123
123
Subsequently, the `AtomicModel` is encapsulated using the `make_model(AtomicModel)` function, which leverages the `deepmd.pt.model.model.make_model.make_model` function. The purpose of the `make_model` wrapper is to facilitate the translation between atomic property predictions and the extended property predictions and differentiation , e.g. the reduction of atomic energy contribution and the autodiff for calculating the forces and virial. The developers usually need to implement an `AtomicModel` not a `Model`.
124
124
@@ -217,6 +217,6 @@ When implementing an existing model in a new backend, directly apply the existin
217
217
218
218
### Consistent tests
219
219
220
-
When transferring features from another backend to the PyTorch backend, it is essential to include a regression test in `/source/tests/consistent` to validate the consistency of the PyTorch backend with other backends. Presently, the regression tests cover self-consistency and cross-backend consistency between TensorFlow, PyTorch, and DP (NumPy) through the serialization/deserialization technique.
220
+
When transferring features from another backend to the PyTorch-TorchScript backend, it is essential to include a regression test in `/source/tests/consistent` to validate the consistency of the PyTorch-TorchScript backend with other backends. Presently, the regression tests cover self-consistency and cross-backend consistency between TensorFlow, PyTorch-TorchScript, and DP (NumPy) through the serialization/deserialization technique.
221
221
222
-
During the development of new components within the PyTorch backend, it is necessary to provide a DP (NumPy) implementation and incorporate corresponding regression tests. For PyTorch components, developers are also required to include a unit test using `torch.jit`.
222
+
During the development of new components within the PyTorch-TorchScript backend, it is necessary to provide a DP (NumPy) implementation and incorporate corresponding regression tests. For PyTorch components, developers are also required to include a unit test using `torch.jit`.
JAX compression accepts `.jax` and `.hlo` inputs. Use `.jax` for the general,
112
+
lossless compressed serialization path. Descriptor pages document whether a
113
+
compressed model can also be exported to StableHLO `.hlo`.
114
+
:::
115
+
116
+
:::{tab-item} DP {{ dpmodel_icon }}
117
+
118
+
```bash
119
+
dp --dp compress -i model.dp -o model-compress.dp
120
+
```
121
+
122
+
DP compression accepts native `.dp`, `.yaml`, and `.yml` models.
123
+
:::
124
+
80
125
::::
81
126
82
-
where `-i` gives the original frozen model, `-o` gives the compressed model. Several other command line options can be passed to `dp compress`, which can be checked with
127
+
where `-i` gives the original frozen model, `-o` gives the compressed model.
128
+
The compression entrypoints resolve the minimum neighbor distance and tabulate
129
+
supported descriptor embedding networks. If the model does not contain a
130
+
minimum neighbor distance, pass the training script with `-t` or
131
+
`--training-script` so it can be computed from the training data.
132
+
133
+
Several other command line options can be passed to `dp compress`, which can be checked with
83
134
84
135
```bash
85
-
$ dp compress --help
136
+
dp compress --help
86
137
```
87
138
88
139
An explanation will be provided
@@ -122,7 +173,7 @@ See the documentation of a specific descriptor to see whether it supports model
122
173
123
174
## Requirements of installation {{ pytorch_icon }}
124
175
125
-
When compressing models in the PyTorch backend, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md).
176
+
When compressing models in the PyTorch-TorchScript backend, the customized OP library for the Python interface must be installed when [freezing the model](../freeze/freeze.md).
126
177
127
178
The customized OP library for the Python interface is installed by default when building DeePMD-kit from source; see the [installation guide](../install/install-from-source.md) for details.
Copy file name to clipboardExpand all lines: doc/freeze/freeze.md
+61-6Lines changed: 61 additions & 6 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -8,17 +8,32 @@ To freeze a model, typically one does
8
8
:::{tab-item} TensorFlow {{ tensorflow_icon }}
9
9
10
10
```bash
11
-
$ dp freeze -o model.pb
11
+
dp freeze -o model.pb
12
12
```
13
13
14
14
in the folder where the model is trained. The output model is called `model.pb`.
15
15
The idea and part of our code are from [Morgan](https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc).
0 commit comments