Skip to content

Commit

Permalink
Merge branch 'devel' into devel-atomic_weight
Browse files Browse the repository at this point in the history
  • Loading branch information
ChiahsinChu authored Dec 18, 2024
2 parents b4899a1 + e8167ce commit 9d0d7d6
Show file tree
Hide file tree
Showing 51 changed files with 4,742 additions and 444 deletions.
4 changes: 2 additions & 2 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ repos:
exclude: ^source/3rdparty
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.8.2
rev: v0.8.3
hooks:
- id: ruff
args: ["--fix"]
Expand Down Expand Up @@ -60,7 +60,7 @@ repos:
- id: blacken-docs
# C++
- repo: https://github.com/pre-commit/mirrors-clang-format
rev: v19.1.4
rev: v19.1.5
hooks:
- id: clang-format
exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$)
Expand Down
18 changes: 13 additions & 5 deletions deepmd/pd/entrypoints/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -230,7 +230,7 @@ def train(
use_pretrain_script: bool = False,
force_load: bool = False,
output: str = "out.json",
):
) -> None:
log.info("Configuration path: %s", input_file)
SummaryPrinter()()
with open(input_file) as fin:
Expand Down Expand Up @@ -321,18 +321,26 @@ def train(
# save min_nbor_dist
if min_nbor_dist is not None:
if not multi_task:
trainer.model.min_nbor_dist = min_nbor_dist
trainer.model.min_nbor_dist = paddle.to_tensor(
min_nbor_dist,
dtype=paddle.float64,
place=DEVICE,
)
else:
for model_item in min_nbor_dist:
trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item]
trainer.model[model_item].min_nbor_dist = paddle.to_tensor(
min_nbor_dist[model_item],
dtype=paddle.float64,
place=DEVICE,
)
trainer.run()


def freeze(
model: str,
output: str = "frozen_model.json",
head: Optional[str] = None,
):
) -> None:
paddle.set_flags(
{
"FLAGS_save_cf_stack_op": 1,
Expand Down Expand Up @@ -383,7 +391,7 @@ def change_bias(
numb_batch: int = 0,
model_branch: Optional[str] = None,
output: Optional[str] = None,
):
) -> None:
if input_file.endswith(".pd"):
old_state_dict = paddle.load(input_file)
model_state_dict = copy.deepcopy(old_state_dict.get("model", old_state_dict))
Expand Down
6 changes: 1 addition & 5 deletions deepmd/pd/loss/ener.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@
TaskLoss,
)
from deepmd.pd.utils import (
decomp,
env,
)
from deepmd.pd.utils.env import (
Expand Down Expand Up @@ -224,10 +223,7 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):

if self.relative_f is not None:
force_label_3 = force_label.reshape([-1, 3])
# norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f
norm_f = (
decomp.norm(force_label_3, axis=1, keepdim=True) + self.relative_f
)
norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f
diff_f_3 = diff_f.reshape([-1, 3])
diff_f_3 = diff_f_3 / norm_f
diff_f = diff_f_3.reshape([-1])
Expand Down
38 changes: 34 additions & 4 deletions deepmd/pd/model/atomic_model/dp_atomic_model.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
# SPDX-License-Identifier: LGPL-3.0-or-later
import copy
import functools
import logging
from typing import (
Expand Down Expand Up @@ -52,7 +51,7 @@ def __init__(
fitting,
type_map: list[str],
**kwargs,
):
) -> None:
super().__init__(type_map, **kwargs)
ntypes = len(type_map)
self.type_map = type_map
Expand Down Expand Up @@ -201,7 +200,7 @@ def serialize(self) -> dict:

@classmethod
def deserialize(cls, data) -> "DPAtomicModel":
data = copy.deepcopy(data)
data = data.copy()
check_version_compatibility(data.pop("@version", 1), 2, 1)
data.pop("@class", None)
data.pop("type", None)
Expand All @@ -212,6 +211,37 @@ def deserialize(cls, data) -> "DPAtomicModel":
obj = super().deserialize(data)
return obj

def enable_compression(
self,
min_nbor_dist: float,
table_extrapolate: float = 5,
table_stride_1: float = 0.01,
table_stride_2: float = 0.1,
check_frequency: int = -1,
) -> None:
"""Call descriptor enable_compression().
Parameters
----------
min_nbor_dist
The nearest distance between atoms
table_extrapolate
The scale of model extrapolation
table_stride_1
The uniform stride of the first table
table_stride_2
The uniform stride of the second table
check_frequency
The overflow check frequency
"""
self.descriptor.enable_compression(
min_nbor_dist,
table_extrapolate,
table_stride_1,
table_stride_2,
check_frequency,
)

def forward_atomic(
self,
extended_coord,
Expand Down Expand Up @@ -278,7 +308,7 @@ def compute_or_load_stat(
self,
sampled_func,
stat_file_path: Optional[DPPath] = None,
):
) -> None:
"""
Compute or load the statistics parameters of the model,
such as mean and standard deviation of descriptors or the energy bias of the fitting net.
Expand Down
6 changes: 6 additions & 0 deletions deepmd/pd/model/descriptor/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,10 @@
from .descriptor import (
DescriptorBlock,
)
from .dpa1 import (
DescrptBlockSeAtten,
DescrptDPA1,
)
from .env_mat import (
prod_env_mat,
)
Expand All @@ -17,6 +21,8 @@
"BaseDescriptor",
"DescriptorBlock",
"DescrptBlockSeA",
"DescrptBlockSeAtten",
"DescrptDPA1",
"DescrptSeA",
"prod_env_mat",
]
Loading

0 comments on commit 9d0d7d6

Please sign in to comment.