deepmodeling / deepmodeling/deepmd-kit

Unsuccessful Acceleration of Compressed Deep Wannier Models in LAMMPS and Python interfaces

Open
#4,755 3 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

### Bug summary

1. In LAMMPS interface, the prediction of Wannier Centroids (``compute dipole all deeptensor/atom dipole.pb``) and the computation of the atomic forces from Wannier Centroid gradient part (``pair_style deepmd ener.pb``) are not successfully accelerated after the compression of Deep Wannier models.
2. In Python interface, the compressed models are successful to accelerate predicting Wannier Centroids (`deepmd.infer.DeepDipole`), but fail to accelerate the Wannier Centroid gradient (`from deepmd.infer.data_modifier import DipoleChargeModifier`).

The ratios of the speed between compressed model and original model are showed in the table:
|| Python Interface | LAMMPS Interface |
|:--------:|:------:|:----------:|
| WC Prediction | $\times$ 2.5 ~ $\times$ 5 | ~1 |
| WC Gradient | ~1 | ~1 |

### DeePMD-kit Version

DeePMD-kit v-2.2.11

### Backend and its version

Tensorflow 2.16.2

### How did you download the software?

pip

### Input Files, Running Commands, Error Log, etc.

A compressed DW model was used to train and freeze a DPLR model, and then compressed the DPLR model. After that, I used the compressed DPLR model to run the DPMD in LAMMPS. However, when I checked the tail of the LAMMPS output file, I found that the compressed DPLR did not accelerate the `Modify` part, which is the most time-consuming part of the entire computation. In order to locate the problem, I tested the speed of WC prediction in both Python and LAMMPS interfaces, and the speed of DW gradient in Python interface.

### 1. Speed test of WC prediction
The script of Python interface:
```python
import time
import glob

from ase import io, Atoms
import numpy as np

from deepmd.infer import DeepDipole

def run(
atoms: Atoms,
dw_model: str,
n_loop: int = 10,
):
dp = DeepDipole(dw_model)
type_map = dp.get_type_map()

atype = np.zeros(len(atoms), dtype=int)
for ii, _atype in enumerate(type_map):
atype[atoms.symbols == _atype] = ii

ts = []
for _ in range(n_loop):
t = time.perf_counter()
dp.eval(
atoms.get_positions().reshape(1, -1),
atoms.get_cell().reshape(1, -1),
atype,
)
ts.append(time.perf_counter() - t)
return np.array(ts)

if __name__ == "__main__":
fnames = glob.glob("../configs/*.xyz")
fnames.sort()
for ii, fname in enumerate(fnames):
atoms = io.read(fname)
out = run(atoms, "../dw_model.pb", 100)
np.savetxt("dw.%03d.out" % ii, out)

out = run(atoms, "../dw_compressed_model.pb", 100)
np.savetxt("dw-compressed.%03d.out" % ii, out)
```
The output with figure shows significant acceleration, especially in large models:
![Image](https://github.com/user-attachments/assets/e0aa7c0f-3533-43b0-972f-a03db277d481)

In LAMMPS, I use the command `compute dipole all deeptensor/atom dw_model.pb` to test the speed of WC prediction, but the compression seems not working:
![Image](https://github.com/user-attachments/assets/8536788e-6ba0-431a-aea5-484bc437183f)

### 2. Speed test of DW gradient
Here is the script of Python interface, and the time of DW gradient part is the difference between the options `eval_fv=True/False` in `DipoleChargeModifier`:
```python
import time

import glob

from ase import io, Atoms
import numpy as np

from deepmd.infer.data_modifier import DipoleChargeModifier

def run(
atoms: Atoms,
dw_model: str,
n_loop: int = 10,
):
dm = DipoleChargeModifier(
dw_model,
model_charge_map=[-8],
sys_charge_map=[6, 1],
ewald_beta=0.4,
ewald_h=0.5,
)
type_map = dm.get_type_map()

atype = np.zeros(len(atoms), dtype=int)
for ii, _atype in enumerate(type_map):
atype[atoms.symbols == _atype] = ii

ts_total = []
ts_ewald = []
for _ in range(n_loop):
t = time.perf_counter()
dm.eval(
atoms.get_positions().reshape(1, -1),
atoms.get_cell().reshape(1, -1),
atype,
eval_fv=True,
)
ts_total.append(time.perf_counter() - t)

t = time.perf_counter()
dm.eval(
atoms.get_positions().reshape(1, -1),
atoms.get_cell().reshape(1, -1),
atype,
eval_fv=False,
)
ts_ewald.append(time.perf_counter() - t)

return np.array(ts_total), np.array(ts_ewald)

if __name__ == "__main__":
fnames = glob.glob("../configs/*.xyz")
fnames.sort()
for ii, fname in enumerate(fnames):
atoms = io.read(fname)
out = run(atoms, "../dw_model.pb", 50)
# print(out[0] - out[1])
np.savetxt(
"dw.%03d.out" % ii, np.transpose([out[0], out[1]]), header="total ewald"
)

out = run(atoms, "../dw_compressed_model.pb", 50)
# print(out[0] - out[1])
np.savetxt(
"dw-compressed.%03d.out" % ii,
np.transpose([out[0], out[1]]),
header="total ewald",
)
```
Output:
![Image](https://github.com/user-attachments/assets/ad2d1520-1a4a-4f95-af88-ec93fc441e9c)
The compression also does not work for accelerating the computation of DW gradient.

### Steps to Reproduce

Please download the file `compressed_dw_test.zip` from the following link, and run the scripts `run.py` in `lmp_dw/`, `py_dw/` and `py_dm/`.

[compressed dw test](https://1drv.ms/u/c/c70cb33f4e38b0db/EWZoohIFIl1MnSoKxwvnb2MBZcHG1x4goR0drIfRYNe2pA?e=dc385W)

### Further Information, Files, and Links

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.