Skip to content

Commit

Permalink
Merge branch 'devel' into mapping
Browse files Browse the repository at this point in the history
  • Loading branch information
CaRoLZhangxy authored Mar 9, 2024
2 parents 91aa410 + a9bcf41 commit 4d4fbf1
Show file tree
Hide file tree
Showing 19 changed files with 199 additions and 108 deletions.
10 changes: 9 additions & 1 deletion deepmd/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,15 @@
_DICT_VAL = TypeVar("_DICT_VAL")
_PRECISION = Literal["default", "float16", "float32", "float64"]
_ACTIVATION = Literal[
"relu", "relu6", "softplus", "sigmoid", "tanh", "gelu", "gelu_tf"
"relu",
"relu6",
"softplus",
"sigmoid",
"tanh",
"gelu",
"gelu_tf",
"none",
"linear",
]
__all__.extend(
[
Expand Down
59 changes: 51 additions & 8 deletions deepmd/dpmodel/utils/network.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
datetime,
)
from typing import (
Callable,
ClassVar,
Dict,
List,
Expand Down Expand Up @@ -309,14 +310,7 @@ def call(self, x: np.ndarray) -> np.ndarray:
"""
if self.w is None or self.activation_function is None:
raise ValueError("w, b, and activation_function must be set")
if self.activation_function == "tanh":
fn = np.tanh
elif self.activation_function.lower() == "none":

def fn(x):
return x
else:
raise NotImplementedError(self.activation_function)
fn = get_activation_fn(self.activation_function)
y = (
np.matmul(x, self.w) + self.b
if self.b is not None
Expand All @@ -332,6 +326,55 @@ def fn(x):
return y


def get_activation_fn(activation_function: str) -> Callable[[np.ndarray], np.ndarray]:
activation_function = activation_function.lower()
if activation_function == "tanh":
return np.tanh
elif activation_function == "relu":

def fn(x):
# https://stackoverflow.com/a/47936476/9567349
return x * (x > 0)

return fn
elif activation_function in ("gelu", "gelu_tf"):

def fn(x):
# generated by GitHub Copilot
return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x**3)))

return fn
elif activation_function == "relu6":

def fn(x):
# generated by GitHub Copilot
return np.minimum(np.maximum(x, 0), 6)

return fn
elif activation_function == "softplus":

def fn(x):
# generated by GitHub Copilot
return np.log(1 + np.exp(x))

return fn
elif activation_function == "sigmoid":

def fn(x):
# generated by GitHub Copilot
return 1 / (1 + np.exp(-x))

return fn
elif activation_function.lower() in ("none", "linear"):

def fn(x):
return x

return fn
else:
raise NotImplementedError(activation_function)


def make_multilayer_network(T_NetworkLayer, ModuleBase):
class NN(ModuleBase):
"""Native representation of a neural network.
Expand Down
3 changes: 2 additions & 1 deletion deepmd/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -798,7 +798,8 @@ def main():
):
deepmd_main = BACKENDS[args.backend]().entry_point_hook
elif args.command is None:
pass
# help message has been printed in parse_args
return
else:
raise RuntimeError(f"unknown command {args.command}")

Expand Down
6 changes: 3 additions & 3 deletions deepmd/pt/entrypoints/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,7 +134,7 @@ def prepare_trainer_input_single(
DpLoaderSet(
validation_systems,
validation_dataset_params["batch_size"],
model_params_single,
model_params_single["type_map"],
)
if validation_systems
else None
Expand All @@ -143,13 +143,13 @@ def prepare_trainer_input_single(
train_data_single = DpLoaderSet(
training_systems,
training_dataset_params["batch_size"],
model_params_single,
model_params_single["type_map"],
)
else:
train_data_single = DpLoaderSet(
training_systems,
training_dataset_params["batch_size"],
model_params_single,
model_params_single["type_map"],
)
return (
train_data_single,
Expand Down
14 changes: 7 additions & 7 deletions deepmd/pt/loss/denoise.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False):
coord_mask = label["coord_mask"]
type_mask = label["type_mask"]

loss = torch.tensor(0.0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)
loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0]
more_loss = {}
if self.has_coord:
if self.mask_loss_coord:
Expand All @@ -66,9 +66,9 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False):
beta=self.beta,
)
else:
coord_loss = torch.tensor(
0.0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
)
coord_loss = torch.zeros(
1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
)[0]
else:
coord_loss = F.smooth_l1_loss(
updated_coord.view(-1, 3),
Expand All @@ -89,9 +89,9 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False):
reduction="mean",
)
else:
token_loss = torch.tensor(
0.0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
)
token_loss = torch.zeros(
1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE
)[0]
else:
token_loss = F.nll_loss(
F.log_softmax(logits.view(-1, self.ntypes - 1), dim=-1),
Expand Down
2 changes: 1 addition & 1 deletion deepmd/pt/loss/ener.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False):
pref_e = self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * coef
pref_f = self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * coef
pref_v = self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * coef
loss = torch.tensor(0.0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)
loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0]
more_loss = {}
# more_loss['log_keys'] = [] # showed when validation on the fly
# more_loss['test_keys'] = [] # showed when doing dp test
Expand Down
2 changes: 1 addition & 1 deletion deepmd/pt/loss/tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ def forward(self, model_pred, label, natoms, learning_rate=0.0, mae=False):
Other losses for display.
"""
del learning_rate, mae
loss = torch.tensor(0.0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)
loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0]
more_loss = {}
if (
self.has_local_weight
Expand Down
14 changes: 8 additions & 6 deletions deepmd/pt/model/atomic_model/linear_atomic_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,10 @@ def __init__(

self.atomic_bias = None
self.mixed_types_list = [model.mixed_types() for model in self.models]
self.rcuts = torch.tensor(
self.get_model_rcuts(), dtype=torch.float64, device=env.DEVICE
)
self.nsels = torch.tensor(self.get_model_nsels(), device=env.DEVICE)
BaseAtomicModel.__init__(self, **kwargs)

def mixed_types(self) -> bool:
Expand Down Expand Up @@ -117,14 +121,12 @@ def get_model_sels(self) -> List[List[int]]:
"""Get the sels for each individual models."""
return [model.get_sel() for model in self.models]

def _sort_rcuts_sels(self, device: torch.device) -> Tuple[List[float], List[int]]:
def _sort_rcuts_sels(self) -> Tuple[List[float], List[int]]:
# sort the pair of rcut and sels in ascending order, first based on sel, then on rcut.
rcuts = torch.tensor(self.get_model_rcuts(), dtype=torch.float64, device=device)
nsels = torch.tensor(self.get_model_nsels(), device=device)
zipped = torch.stack(
[
rcuts,
nsels,
self.rcuts,
self.nsels,
],
dim=0,
).T
Expand Down Expand Up @@ -171,7 +173,7 @@ def forward_atomic(
if self.do_grad_r() or self.do_grad_c():
extended_coord.requires_grad_(True)
extended_coord = extended_coord.view(nframes, -1, 3)
sorted_rcuts, sorted_sels = self._sort_rcuts_sels(device=extended_coord.device)
sorted_rcuts, sorted_sels = self._sort_rcuts_sels()
nlists = build_multiple_neighbor_list(
extended_coord,
nlist,
Expand Down
10 changes: 6 additions & 4 deletions deepmd/pt/train/training.py
Original file line number Diff line number Diff line change
Expand Up @@ -973,9 +973,11 @@ def get_data(self, is_train=True, task_key="Default"):
continue
elif not isinstance(batch_data[key], list):
if batch_data[key] is not None:
batch_data[key] = batch_data[key].to(DEVICE)
batch_data[key] = batch_data[key].to(DEVICE, non_blocking=True)
else:
batch_data[key] = [item.to(DEVICE) for item in batch_data[key]]
batch_data[key] = [
item.to(DEVICE, non_blocking=True) for item in batch_data[key]
]
# we may need a better way to classify which are inputs and which are labels
# now wrapper only supports the following inputs:
input_keys = [
Expand Down Expand Up @@ -1035,7 +1037,7 @@ def print_on_training(self, fout, step_id, cur_lr, train_results, valid_results)
print_str = ""
print_str += "%7d" % step_id
if not self.multi_task:
if valid_results is not None:
if valid_results:
prop_fmt = " %11.2e %11.2e"
for k in train_keys:
print_str += prop_fmt % (valid_results[k], train_results[k])
Expand All @@ -1045,7 +1047,7 @@ def print_on_training(self, fout, step_id, cur_lr, train_results, valid_results)
print_str += prop_fmt % (train_results[k])
else:
for model_key in self.model_keys:
if valid_results[model_key] is not None:
if valid_results[model_key]:
prop_fmt = " %11.2e %11.2e"
for k in sorted(valid_results[model_key].keys()):
print_str += prop_fmt % (
Expand Down
21 changes: 17 additions & 4 deletions deepmd/pt/utils/dataloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,13 +55,27 @@ def setup_seed(seed):


class DpLoaderSet(Dataset):
"""A dataset for storing DataLoaders to multiple Systems."""
"""A dataset for storing DataLoaders to multiple Systems.
Parameters
----------
sys_path
Path to the data system
batch_size
Max frame count in a batch.
type_map
Gives the name of different atom types
seed
Random seed for dataloader
shuffle
If the data are shuffled (Only effective in serial mode. Always shuffle in distributed data parallelism)
"""

def __init__(
self,
systems,
batch_size,
model_params,
type_map,
seed=10,
shuffle=True,
):
Expand All @@ -77,8 +91,7 @@ def __init__(
def construct_dataset(system):
return DeepmdDataSetForLoader(
system=system,
type_map=model_params["type_map"],
shuffle=shuffle,
type_map=type_map,
)

with Pool(
Expand Down
13 changes: 3 additions & 10 deletions deepmd/pt/utils/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

from typing import (
List,
Optional,
)

from torch.utils.data import (
Expand All @@ -16,24 +17,16 @@


class DeepmdDataSetForLoader(Dataset):
def __init__(
self,
system: str,
type_map: str,
shuffle=True,
):
def __init__(self, system: str, type_map: Optional[List[str]] = None):
"""Construct DeePMD-style dataset containing frames cross different systems.
Args:
- systems: Paths to systems.
- batch_size: Max frame count in a batch.
- type_map: Atom types.
"""
self.system = system
self._type_map = type_map
self._data_system = DeepmdData(
sys_path=system, shuffle_test=shuffle, type_map=self._type_map
)
self._data_system = DeepmdData(sys_path=system, type_map=self._type_map)
self.mixed_type = self._data_system.mixed_type
self._ntypes = self._data_system.get_ntypes()
self._natoms = self._data_system.get_natoms()
Expand Down
20 changes: 16 additions & 4 deletions deepmd/pt/utils/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,16 @@ def get_activation_fn(activation: str) -> Callable:
"""Returns the activation function corresponding to `activation`."""
if activation.lower() == "relu":
return F.relu
elif activation.lower() == "gelu":
return F.gelu
elif activation.lower() == "gelu" or activation.lower() == "gelu_tf":
return lambda x: F.gelu(x, approximate="tanh")
elif activation.lower() == "tanh":
return torch.tanh
elif activation.lower() == "relu6":
return F.relu6
elif activation.lower() == "softplus":
return F.softplus
elif activation.lower() == "sigmoid":
return torch.sigmoid
elif activation.lower() == "linear" or activation.lower() == "none":
return lambda x: x
else:
Expand All @@ -42,10 +48,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:

if self.activation.lower() == "relu":
return F.relu(x)
elif self.activation.lower() == "gelu":
return F.gelu(x)
elif self.activation.lower() == "gelu" or self.activation.lower() == "gelu_tf":
return F.gelu(x, approximate="tanh")
elif self.activation.lower() == "tanh":
return torch.tanh(x)
elif self.activation.lower() == "relu6":
return F.relu6(x)
elif self.activation.lower() == "softplus":
return F.softplus(x)
elif self.activation.lower() == "sigmoid":
return torch.sigmoid(x)
elif self.activation.lower() == "linear" or self.activation.lower() == "none":
return x
else:
Expand Down
Loading

0 comments on commit 4d4fbf1

Please sign in to comment.