Skip to content

Commit

Permalink
polish code
Browse files Browse the repository at this point in the history
  • Loading branch information
@HydrogenSulfate
HydrogenSulfate committed Nov 28, 2023
1 parent 46dbc9c commit b5b5d54
Show file tree
Hide file tree
Showing 8 changed files with 310 additions and 418 deletions.
362 changes: 170 additions & 192 deletions deepmd/descriptor/se_a.py
View file

Large diffs are not rendered by default.

71 changes: 16 additions & 55 deletions deepmd/entrypoints/freeze.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -338,40 +338,6 @@ def freeze_graph(
out_suffix : str
The chosen suffix to freeze in the input_graph.
"""
# output_node = _make_node_names(
# freeze_type, modifier, out_suffix=out_suffix, node_names=node_names
# )
# different_set = set(output_node) - set(input_node)
# if different_set:
# log.warning(
# "The following nodes are not in the graph: %s. "
# "Skip freezeing these nodes. You may be freezing "
# "a checkpoint generated by an old version." % different_set
# )
# # use intersection as output list
# output_node = list(set(output_node) & set(input_node))
# log.info(f"The following nodes will be frozen: {output_node}")
# # We use a built-in TF helper to export variables to constants
# output_graph_def = tf.graph_util.convert_variables_to_constants(
# sess, # The session is used to retrieve the weights
# input_graph, # The graph_def is used to retrieve the nodes
# output_node, # The output node names are used to select the usefull nodes
# )
# # if multi-task, change fitting_net suffix and model_type
# if out_suffix != "":
# output_graph_def = _modify_model_suffix(
# output_graph_def, out_suffix, freeze_type
# )

# # If we need to transfer the fitting net variables
# output_graph_def = _transfer_fitting_net_trainable_variables(
# sess, output_graph_def, input_graph
# )

# # Finally we serialize and dump the output graph to the filesystem
# with tf.gfile.GFile(out_graph_name, "wb") as f:
# f.write(output_graph_def.SerializeToString())
# log.info(f"{len(output_graph_def.node):d} ops in the final graph.")
import paddle

from deepmd.infer import DeepPot
Expand All @@ -381,15 +347,6 @@ def freeze_graph(
load_prefix="load",
default_tf_graph=False,
)
# print(dp.model.descrpt.embedding_nets[0][0].weight[0])
# for w in dp.model.descrpt.embedding_nets[0][0].weight:
# print(f"w {w.shape} {w.mean().item()} {w.var().item()}")
# print("从state_dict打印载入的参数")
# for k, v in dp.model.state_dict().items():
# print(f"{k} {v.shape} {v.dtype} {v.mean().item()} {v.var().item()}")
# exit()
# for b in dp.model.descrpt.embedding_nets[0][0].bias:
# print(f"b {b.shape} {b.mean().item()} {b.var().item()}")
dp.model.eval()
from paddle.static import InputSpec

Expand Down Expand Up @@ -424,21 +381,25 @@ def freeze_graph(
False,
],
)
print(f"st_model.descrpt.buffer_rcut.name = {st_model.descrpt.buffer_rcut.name}")
print(
f"st_model.descrpt.buffer_ntypes.name = {st_model.descrpt.buffer_ntypes.name}"
)
print(
f"st_model.fitting.buffer_dfparam.name = {st_model.fitting.buffer_dfparam.name}"
)
print(
f"st_model.fitting.buffer_daparam.name = {st_model.fitting.buffer_daparam.name}"
)
for name, param in st_model.named_buffers():
print(
f"[{name}, {param.shape}] generated name in static_model is: {param.name}"
)
# print(f"st_model.descrpt.buffer_rcut.name = {st_model.descrpt.buffer_rcut.name}")
# print(
# f"st_model.descrpt.buffer_ntypes.name = {st_model.descrpt.buffer_ntypes.name}"
# )
# print(
# f"st_model.fitting.buffer_dfparam.name = {st_model.fitting.buffer_dfparam.name}"
# )
# print(
# f"st_model.fitting.buffer_daparam.name = {st_model.fitting.buffer_daparam.name}"
# )
# 跳过对program的裁剪，从而保留rcut、ntypes等不参与前向的参数，从而在C++端可以获取这些参数
skip_prune_program = True
print(f"==>> skip_prune_program = {skip_prune_program}")
print(f"==>> Set skip_prune_program = {skip_prune_program}")
paddle.jit.save(st_model, output, skip_prune_program=skip_prune_program)
print(f"Saved to path: {output}")
print(f"Infernece model has been saved to: {output}")


def freeze_graph_multi(
Expand Down
29 changes: 10 additions & 19 deletions deepmd/fit/ener.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ def __init__(
"""Constructor."""
# model param
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
self.dim_descrpt = descrpt.get_dim_out() # M1*M2
self.use_aparam_as_mask = use_aparam_as_mask
# args = ()\
# .add('numb_fparam', int, default = 0)\
Expand Down Expand Up @@ -255,11 +255,6 @@ def __init__(
)
)

# print("create bias_atom_e", self.bias_atom_e.shape, self.bias_atom_e)
# self.register_buffer(
# "t_bias_atom_e",
# paddle.to_tensor(self.bias_atom_e),
# )
if self.numb_fparam > 0:
if self.fparam_avg is None:
self.fparam_avg = 0.0
Expand Down Expand Up @@ -316,10 +311,6 @@ def compute_output_stats(self, all_stat: dict, mixed_type: bool = False) -> None
all_stat, rcond=self.rcond, mixed_type=mixed_type
)
paddle.assign(self.bias_atom_e, self.t_bias_atom_e)
# self.register_buffer(
# "t_bias_atom_e",
# paddle.to_tensor(self.bias_atom_e),
# )

def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
data = all_stat["energy"]
Expand Down Expand Up @@ -439,7 +430,7 @@ def _build_lower(
[0, start_index, 0],
[inputs.shape[0], start_index + natoms, inputs.shape[2]],
)
inputs_i = paddle.reshape(inputs_i, [-1, self.dim_descrpt])
inputs_i = paddle.reshape(inputs_i, [-1, self.dim_descrpt]) # [natoms, M1*M2]
layer = inputs_i
if fparam is not None:
ext_fparam = paddle.tile(fparam, [1, natoms])
Expand Down Expand Up @@ -504,7 +495,7 @@ def _build_lower(
if (not self.uniform_seed) and (self.seed is not None):
self.seed += self.seed_shift

return final_layer
return final_layer # [natoms, 1]

def forward(
self,
Expand All @@ -519,7 +510,7 @@ def forward(
Parameters
----------
inputs
The input descriptor
The input descriptor, [1, all_atoms, M1*M2]
input_dict
Additional dict for inputs.
if numb_fparam > 0, should have input_dict['fparam']
Expand Down Expand Up @@ -575,7 +566,9 @@ def forward(
self.bias_atom_e[type_i] = self.bias_atom_e[type_i]
self.bias_atom_e = self.bias_atom_e[:ntypes_atom]

inputs = paddle.reshape(inputs, [-1, natoms[0], self.dim_descrpt])
inputs = paddle.reshape(
inputs, [-1, natoms[0], self.dim_descrpt]
) # [1, all_atoms, M1*M2]
if len(self.atom_ener):
# only for atom_ener
nframes = input_dict.get("nframes")
Expand Down Expand Up @@ -680,7 +673,7 @@ def forward(
final_layer -= zero_layer
final_layer = paddle.reshape(
final_layer, [paddle.shape(inputs)[0], natoms[2 + type_i]]
)
) # [1, natoms]
outs_list.append(final_layer)
start_index += natoms[2 + type_i]
# concat the results
Expand Down Expand Up @@ -727,9 +720,7 @@ def forward(
),
[paddle.shape(inputs)[0], paddle.sum(natoms[2 : 2 + ntypes_atom]).item()],
)
# print(__file__, self.t_bias_atom_e)
# exit()
outs = outs + self.add_type
outs = outs + self.add_type # 类型编码(类似于transformer的位置编码，每种类型自己有一个特征，加到原特征上)
outs *= atype_filter
self.atom_ener_after = outs

Expand All @@ -747,7 +738,7 @@ def forward(
)
outs = outs - outs_mean
outs = paddle.reshape(outs, [-1])
return paddle.reshape(outs, [-1])
return paddle.reshape(outs, [-1]) # [all_atoms]

def init_variables(
self,
Expand Down
63 changes: 40 additions & 23 deletions deepmd/infer/deep_eval.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from deepmd.common import expand_sys_str
from deepmd.common import j_loader
from deepmd.common import j_must_have
from deepmd.descriptor import DescrptSeA
from deepmd.env import MODEL_VERSION
from deepmd.env import default_tf_session_config
from deepmd.env import paddle
Expand Down Expand Up @@ -54,12 +55,24 @@ def __init__(
model_param = j_must_have(jdata, "model")

descrpt_param = j_must_have(model_param, "descriptor")
from deepmd.descriptor import DescrptSeA
explicit_ntypes_descrpt = ["se_atten"]
# hybrid_with_tebd = False
if descrpt_param["type"] in explicit_ntypes_descrpt:
descrpt_param["ntypes"] = len(model_param["type_map"])
elif descrpt_param["type"] == "hybrid":
for descrpt_item in descrpt_param["list"]:
if descrpt_item["type"] in explicit_ntypes_descrpt:
descrpt_item["ntypes"] = len(model_param["type_map"])
# hybrid_with_tebd = True

# if descrpt_param["type"] in ["se_e2_a", "se_a", "se_e2_r", "se_r", "hybrid"]:
descrpt_param["spin"] = None
descrpt_param["type_one_side"] = False

descrpt_param.pop("type", None)
descrpt_param.pop("_comment", None)
self.spin = None
descrpt_param["spin"] = self.spin
# descrpt_param["spin"] = self.spin
self.descrpt = DescrptSeA(**descrpt_param)

self.multi_task_mode = "fitting_net_dict" in model_param
Expand Down Expand Up @@ -90,28 +103,32 @@ def __init__(
model_param.get("sw_rmax"),
self.spin,
)
load_state_dict = paddle.load(str(model_file))
for k, v in load_state_dict.items():
if k in self.model.state_dict():
if load_state_dict[k].dtype != self.model.state_dict()[k].dtype:
print(
f"convert {k}'s dtype from {load_state_dict[k].dtype} to {self.model.state_dict()[k].dtype}"
)
load_state_dict[k] = load_state_dict[k].astype(
self.model.state_dict()[k].dtype
)
if list(load_state_dict[k].shape) != list(
self.model.state_dict()[k].shape
):
print(
f"convert {k}'s shape from {load_state_dict[k].shape} to {self.model.state_dict()[k].shape}"
)
load_state_dict[k] = load_state_dict[k].reshape(
model_file_str = str(model_file)
if model_file_str.endswith((".pdmodel", ".pdiparams")):
st_model_prefix = model_file_str.rsplit(".", 1)[0]
self.st_model = paddle.jit.load(st_model_prefix)
else:
load_state_dict = paddle.load(str(model_file))
for k, v in load_state_dict.items():
if k in self.model.state_dict():
if load_state_dict[k].dtype != self.model.state_dict()[k].dtype:
print(
f"convert {k}'s dtype from {load_state_dict[k].dtype} to {self.model.state_dict()[k].dtype}"
)
load_state_dict[k] = load_state_dict[k].astype(
self.model.state_dict()[k].dtype
)
if list(load_state_dict[k].shape) != list(
self.model.state_dict()[k].shape
)
# print(f"==>> Load pretraied model successfully from: {str(model_file)}")
# exit()
self.model.set_state_dict(load_state_dict)
):
print(
f"convert {k}'s shape from {load_state_dict[k].shape} to {self.model.state_dict()[k].shape}"
)
load_state_dict[k] = load_state_dict[k].reshape(
self.model.state_dict()[k].shape
)
self.model.set_state_dict(load_state_dict)
print(f"==>> Load pretraied model successfully from: {str(model_file)}")
self.load_prefix = load_prefix

# graph_compatable should be called after graph and prefix are set
Expand Down
70 changes: 2 additions & 68 deletions deepmd/infer/deep_pot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -202,40 +202,6 @@ def __init__(
# ewald_beta=ewald_beta,
# )

# NOTE: 使用静态图模型推理
if not hasattr(self, "st_model"):
self.st_model = paddle.jit.load(
"/workspace/hesensen/deepmd_backend/deepmd-kit/examples/water/se_e2_a/Model_1000000"
)
# for k, v in self.st_model.named_parameters():
# print(f"{k} {v.shape} {v.mean().item()} {v.var().item()}")
# """
# param_0 [1, 25] 0.9498768667019655 0.7340928425051493
# param_1 [1, 50] 1.1214760345730344 0.9621536430386503
# param_2 [1, 100] 1.168418946306086 1.0411743399117217
# param_3 [1, 25] 0.002546645920014433 0.27806176560439083
# param_4 [25, 50] -0.015372691466039676 0.10679961485782502
# param_5 [50, 100] -0.0010681208730640539 0.09950205346985407
# param_6 [1, 25] 1.0639599744616117 0.917256936729768
# param_7 [1, 50] 1.142691803888668 0.9639366693005659
# param_8 [1, 100] 1.1471394365452061 1.0091294911290036
# param_9 [1, 25] 0.019013792716200625 0.1450311660373793
# param_10 [25, 50] -0.006747145320748169 0.028971429954693633
# param_11 [50, 100] -0.03750622755877242 0.04714041793007081
# param_12 [1, 25] 1.0380588819220322 0.8904020425094114
# param_13 [1, 50] 1.1245407895732316 0.9234643810098301
# param_14 [1, 100] 1.1430567514092813 0.9876968977916372
# param_15 [1, 25] 0.03272738992064966 0.1751917732380509
# param_16 [25, 50] -0.017871745658352124 0.0384813911462805
# param_17 [50, 100] -0.07345191324160481 0.1768254187693918
# param_18 [1, 25] 1.0147830400771964 0.9070964180637516
# param_19 [1, 50] 1.1198266551333698 1.034746190888665
# param_20 [1, 100] 1.1410748813679754 1.0428001731414345
# param_21 [1, 25] -0.022862385119536602 0.18038150422614693
# param_22 [25, 50] -0.024970130750642985 0.07176423978220656
# param_23 [50, 100] -0.012309303874398866 0.07227932085917015
# """

def _run_default_sess(self):
if self.has_spin is True:
[
Expand Down Expand Up @@ -611,7 +577,6 @@ def _eval_inner(
# print(eval_inputs['type'].shape) # [960]
# print(eval_inputs['natoms_vec'].shape) # [4]
# print(eval_inputs['box'].shape) # [45]
# exit()

if self.has_fparam:
eval_inputs["fparam"] = paddle.to_tensor(
Expand All @@ -629,24 +594,14 @@ def _eval_inner(
# eval_inputs['default_mesh'] = paddle.to_tensor(np.array([], dtype = np.int32))

if hasattr(self, "st_model"):
# NOTE: 使用静态图模型推理
eval_outputs = self.st_model(
eval_inputs["coord"], # [2880] paddle.float64
eval_inputs["type"], # [960] paddle.int32
eval_inputs["natoms_vec"], # [4] paddle.int32
eval_inputs["box"], # [45] paddle.float64
eval_inputs["default_mesh"], # [6] paddle.int32
)
# print(eval_inputs["coord"].shape)
# print(eval_inputs["type"].shape)
# print(eval_inputs["natoms_vec"].shape)
# print(eval_inputs["box"].shape)
# print(eval_inputs["default_mesh"].shape)
# np.save("/workspace/hesensen/deepmd_backend/python_infer_data/coord.npy", eval_inputs["coord"].numpy())
# np.save("/workspace/hesensen/deepmd_backend/python_infer_data/type.npy", eval_inputs["type"].numpy())
# np.save("/workspace/hesensen/deepmd_backend/python_infer_data/natoms_vec.npy", eval_inputs["natoms_vec"].numpy())
# np.save("/workspace/hesensen/deepmd_backend/python_infer_data/box.npy", eval_inputs["box"].numpy())
# np.save("/workspace/hesensen/deepmd_backend/python_infer_data/default_mesh.npy", eval_inputs["default_mesh"].numpy())
# exit()
eval_outputs = {
"atom_ener": eval_outputs[0],
"atom_virial": eval_outputs[1],
Expand All @@ -655,37 +610,16 @@ def _eval_inner(
"energy": eval_outputs[4],
"force": eval_outputs[5],
"virial": eval_outputs[6],
# "z00_hidden1": eval_outputs[7],
# "z00_hidden2": eval_outputs[8],
# "z00_hidden3": eval_outputs[9],
# "z00_xx1": eval_outputs[7],
# "z00_xx2": eval_outputs[8],
# "z00_xx3": eval_outputs[9],
# "z00_xx4": eval_outputs[10],
# "weight_0": eval_outputs[7],
# "bias_0": eval_outputs[8],
# "xx1": eval_outputs[9],
# "hidden1": eval_outputs[10],
}

# for k, v in eval_outputs.items():
# print(k, v.shape)
# np.save(f"/workspace/hesensen/deepmd_backend/python_infer_data/st_model_{k}.npy", v.numpy())
# print(f"finished save {k}")
# exit()
else:
# NOTE: 使用动态图模型推理
eval_outputs = self.model(
eval_inputs["coord"], # [2880] paddle.float64
eval_inputs["type"], # [960] paddle.int32
eval_inputs["natoms_vec"], # [4] paddle.int32
eval_inputs["box"], # [45] paddle.float64
eval_inputs["default_mesh"], # [6] paddle.int32
eval_inputs,
# eval_inputs.coord: [2880] paddle.float64
# eval_inputs.type: [960] paddle.int32
# eval_inputs.natoms_vec: [4] paddle.int32
# eval_inputs.box: [45] paddle.float64
# eval_inputs.default_mesh: [6] paddle.int32
suffix="",
reuse=False,
)
Expand Down
Loading

0 comments on commit b5b5d54

Please sign in to comment.