deepmodeling / deepmodeling/deepmd-kit
Unsuccessful Acceleration of Compressed Deep Wannier Models in LAMMPS and Python interfaces
- 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:

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:

### 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:

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
Assessment
This issue has not been assessed yet.