compress.md

Compress a model {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }}

Note

Backends covered below: TensorFlow and TensorFlow 2 {{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable {{ pytorch_icon }}, JAX {{ jax_icon }}, and DP {{ dpmodel_icon }}.

Theory

The compression of the DP model uses three techniques, tabulated inference, operator merging, and precise neighbor indexing, to improve the performance of model training and inference when the model parameters are properly trained.

For better performance, the NN inference can be replaced by tabulated function evaluations if the input of the NN is of dimension one. The idea is to approximate the output of the NN by a piece-wise polynomial fitting. The input domain (a compact domain in $\mathbb R$) is divided into $L_c$ equally spaced intervals, in which we apply a fifth-order polynomial $g^l_m(x)$ approximation of the $m$-th output component of the NN function:

$$g^l_m(x) = a^l_m x^5 + b^l_m x^4 + c^l_m x^3 + d^l_m x^2 + e^l_m x + f^l_m,\quad x \in [x_l, x_{l+1}),$$

where $l=1,2,\dots,L_c$ is the index of the intervals, $x_1, \dots, x_{L_c}, x_{L_c+1}$ are the endpoints of the intervals, and $a^l_m$, $b^l_m$, $c^l_m$, $d^l_m$, $e^l_m$, and $f^l_m$ are the fitting parameters. The fitting parameters can be computed by the equations below:

$$a^l_m = \frac{1}{2\Delta x_l^5}[12h_{m,l}-6(y'_{m,l+1}+y'_{m,l})\Delta x_l + (y''_{m,l+1}-y''_{m,l})\Delta x_l^2],$$ $$b^l_m = \frac{1}{2\Delta x_l^4}[-30h_{m,l} +(14y'_{m,l+1}+16y'_{m,l})\Delta x_l + (-2y''_{m,l+1}+3y''_{m,l})\Delta x_l^2],$$ $$c^l_m = \frac{1}{2\Delta x_l^3}[20h_{m,l}-(8y'_{m,l+1}+12y'_{m,l})\Delta x_l + (y''_{m,l+1}-3y''_{m,l})\Delta x_l^2],$$ $$d^l_m = \frac{1}{2}y''_{m,l},$$ $$e^l_m = y_{m,l}',$$ $$f^l_m = y_{m,l},$$

where $\Delta x_l=x_{l+1}-x_l$ denotes the size of the interval. $h_{m,l}=y_{m,l+1}-y_{m,l}$. $y_{m,l} = y_m(x_l)$, $y'_{m,l} = y'm(x_l)$ and $y''{m,l} = y''_m(x_l)$ are the value, the first-order derivative, and the second-order derivative of the $m$-th component of the target NN function at the interval point $x_l$, respectively. The first and second-order derivatives are easily calculated by the back-propagation of the NN functions.

In the standard DP model inference, taking the two-body embedding descriptor as an example, the matrix product $(\mathcal G^i)^T \mathcal R$ requires the transfer of the tensor $\mathcal G^i$ between the register and the host/device memories, which usually becomes the bottle-neck of the computation due to the relatively small memory bandwidth of the GPUs. The compressed DP model merges the matrix multiplication $(\mathcal G^i)^T \mathcal R$ with the tabulated inference step. More specifically, once one column of the $(\mathcal G^i)^T$ is evaluated, it is immediately multiplied with one row of the environment matrix in the register, and the outer product is deposited to the result of $(\mathcal G^i)^T \mathcal R$. By the operator merging technique, the allocation of $\mathcal G^i$ and the memory movement between register and host/device memories is avoided. The operator merging of the three-body embedding can be derived analogously.

The first dimension, $N_c$, of the environment ($\mathcal R^i$) and embedding ($\mathcal G^i$) matrices is the expected maximum number of neighbors. If the number of neighbors of an atom is smaller than $N_c$, the corresponding positions of the matrices are pad with zeros. In practice, if the real number of neighbors is significantly smaller than $N_c$, a notable operation is spent on the multiplication of padding zeros. In the compressed DP model, the number of neighbors is precisely indexed at the tabulated inference stage, further saving computational costs.¹

Instructions

Once the frozen model is obtained from DeePMD-kit, we can get the neural network structure and its parameters (weights, biases, etc.) from the trained model, and compress it in the following way:

::::{tab-set}

:::{tab-item} TensorFlow {{ tensorflow_icon }}

dp compress -i graph.pb -o graph-compress.pb

:::

:::{tab-item} TensorFlow 2 {{ tensorflow_icon }}

dp --tf2 compress -i model.ckpt.tf2 -o model-compress.savedmodeltf

TensorFlow 2 compression reads a .tf2 training checkpoint directory or a checkpoint prefix and writes a compressed .savedmodeltf model. See the descriptor documentation for model-specific SavedModel export requirements. :::

:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }}

dp --pt compress -i model.pth -o model-compress.pth

:::

:::{tab-item} PyTorch-Exportable {{ pytorch_icon }}

dp --pt-expt compress -i dpa1-graph.pt2 -o dpa1-graph-compress.pt2

This command produces an executable compressed artifact only for descriptors that support the PyTorch-Exportable compression path. See the DPA-1 model-compression requirements for the graph-lowered .pt2 route used in this example. :::

:::{tab-item} JAX {{ jax_icon }}

dp --jax compress -i frozen_model.jax -o compressed_model.jax

JAX compression accepts .jax and .hlo inputs. Use .jax for the general, lossless compressed serialization path. Descriptor pages document whether a compressed model can also be exported to StableHLO .hlo. :::

:::{tab-item} DP {{ dpmodel_icon }}

dp --dp compress -i model.dp -o model-compress.dp

DP compression accepts native .dp, .yaml, and .yml models. :::

::::

where -i gives the original frozen model, -o gives the compressed model. The compression entrypoints resolve the minimum neighbor distance and tabulate supported descriptor embedding networks. If the model does not contain a minimum neighbor distance, pass the training script with -t or --training-script so it can be computed from the training data.

Several other command line options can be passed to dp compress, which can be checked with

dp compress --help

An explanation will be provided

Parameter explanation

Model compression, which includes tabulating the embedding net. The table is composed of fifth-order polynomial coefficients and is assembled from two sub-tables. For model descriptor with se_e2_a type, the first sub-table takes the stride(parameter) as its uniform stride, while the second sub-table takes 10 _ stride as its uniform stride; For model descriptor with se_e3 type, the first sub-table takes 10 _ stride as it's uniform stride, while the second sub-table takes 100 _ stride as it's uniform stride. The range of the first table is automatically detected by DeePMD-kit, while the second table ranges from the first table's upper boundary(upper) to the extrapolate(parameter) _ upper. Finally, we added a check frequency parameter. It indicates how often the program checks for overflow(if the input environment matrix overflows the first or second table range) during the MD inference.

Justification of model compression

Model compression, with little loss of accuracy, can greatly speed up MD inference time. According to different simulation systems and training parameters, the speedup can reach more than 10 times at both CPU and GPU devices. At the same time, model compression can greatly change memory usage, reducing as much as 20 times under the same hardware conditions.

Acceptable original model version

The model compression interface requires the version of DeePMD-kit used in the original model generation should be 2.0.0-alpha.0 or above. If one has a frozen 1.2 or 1.3 model, one can upgrade it through the dp convert-from interface. (eg: dp convert-from 1.2/1.3 -i old_frozen_model.pb -o new_frozen_model.pb)

Acceptable descriptor type

Not any descriptor supports model compression. See the documentation of a specific descriptor to see whether it supports model compression.

Available activation functions for descriptor:

• tanh
• gelu
• relu
• relu6
• softplus
• sigmoid

Requirements of installation {{ pytorch_icon }}

When compressing models in the PyTorch-TorchScript backend, the customized OP library for the Python interface must be installed when freezing the model.

The customized OP library for the Python interface is installed by default when building DeePMD-kit from source; see the installation guide for details.

Footnotes

1. This section is built upon Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang, J. Chem. Phys. 159, 054801 (2023) licensed under a Creative Commons Attribution (CC BY) license. ↩