Skip to content

Commit

Permalink
update reformatting code
Browse files Browse the repository at this point in the history
  • Loading branch information
@HydrogenSulfate
HydrogenSulfate committed Dec 22, 2023
2 parents 3eb254c + 36e0082 commit 52638f2
Show file tree
Hide file tree
Showing 13 changed files with 84 additions and 770 deletions.
217 changes: 39 additions & 178 deletions deepmd/descriptor/se_a.py
View file

Large diffs are not rendered by default.

85 changes: 9 additions & 76 deletions deepmd/entrypoints/freeze.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -311,6 +311,14 @@ def _make_node_names(
def freeze_graph(
model_file: str,
output: str,
# sess,
# input_graph,
# input_node,
# freeze_type,
# modifier,
# out_graph_name,
# node_names=None,
# out_suffix="",
):
"""Freeze the single graph with chosen out_suffix.
Expand Down Expand Up @@ -356,7 +364,7 @@ def freeze_graph(
print(
f"[{name}, {param.dtype}, {param.shape}] generated name in static_model is: {param.name}"
)
# 跳过对program的裁剪，从而保留rcut、ntypes等不参与前向的参数，从而在C++端可以获取这些参数
# skip pruning for program so as to keep buffers into files
skip_prune_program = True
print(f"==>> Set skip_prune_program = {skip_prune_program}")
paddle.jit.save(st_model, output, skip_prune_program=skip_prune_program)
Expand Down Expand Up @@ -441,12 +449,8 @@ def freeze_graph_multi(

def freeze(
*,
# checkpoint_folder: str,
input_file: str,
output: str,
# node_names: Optional[str] = None,
# nvnmd_weight: Optional[str] = None,
# united_model: bool = False,
**kwargs,
):
"""Freeze the graph in supplied folder.
Expand All @@ -460,78 +464,7 @@ def freeze(
**kwargs
other arguments
"""
# We retrieve our checkpoint fullpath
# checkpoint = tf.train.get_checkpoint_state(checkpoint_folder)
# input_checkpoint = checkpoint.model_checkpoint_path

# # expand the output file to full path
# output_graph = abspath(output)

# # Before exporting our graph, we need to precise what is our output node
# # This is how TF decides what part of the Graph he has to keep
# # and what part it can dump
# # NOTE: this variable is plural, because you can have multiple output nodes
# # node_names = "energy_test,force_test,virial_test,t_rcut"

# # We clear devices to allow TensorFlow to control
# # on which device it will load operations
# clear_devices = True

# # We import the meta graph and retrieve a Saver
# try:
# # In case paralle training
# import horovod.tensorflow as _ # noqa: F401
# except ImportError:
# pass
# saver = tf.train.import_meta_graph(
# f"{input_checkpoint}.meta", clear_devices=clear_devices
# )

# # We retrieve the protobuf graph definition
# graph = tf.get_default_graph()
# try:
# input_graph_def = graph.as_graph_def()
# except google.protobuf.message.DecodeError as e:
# raise GraphTooLargeError(
# "The graph size exceeds 2 GB, the hard limitation of protobuf."
# " Then a DecodeError was raised by protobuf. You should "
# "reduce the size of your model."
# ) from e
# nodes = [n.name for n in input_graph_def.node]

# # We start a session and restore the graph weights
# with tf.Session() as sess:
# saver.restore(sess, input_checkpoint)
# model_type = run_sess(sess, "model_attr/model_type:0", feed_dict={}).decode(
# "utf-8"
# )
# if "modifier_attr/type" in nodes:
# modifier_type = run_sess(sess, "modifier_attr/type:0", feed_dict={}).decode(
# "utf-8"
# )
# else:
# modifier_type = None
# if nvnmd_weight is not None:
# save_weight(sess, nvnmd_weight) # nvnmd
# if model_type != "multi_task":
freeze_graph(
input_file,
output,
# sess,
# input_graph_def,
# nodes,
# model_type,
# modifier_type,
# output_graph,
# node_names,
)
# else:
# freeze_graph_multi(
# sess,
# input_graph_def,
# nodes,
# modifier_type,
# output_graph,
# node_names,
# united_model=united_model,
# )
13 changes: 1 addition & 12 deletions deepmd/entrypoints/train.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -419,15 +419,6 @@ def get_nbor_stat(jdata, rcut, one_type: bool = False):

min_nbor_dist, max_nbor_size = neistat.get_stat(train_data)

# moved from traier.py as duplicated
# TODO: this is a simple fix but we should have a clear
# architecture to call neighbor stat
# tf.constant(
# min_nbor_dist,
# name="train_attr/min_nbor_dist",
# dtype=GLOBAL_ENER_FLOAT_PRECISION,
# )
# tf.constant(max_nbor_size, name="train_attr/max_nbor_size", dtype=tf.int32)
return min_nbor_dist, max_nbor_size


Expand Down Expand Up @@ -473,9 +464,7 @@ def update_one_sel(jdata, descriptor):
if descriptor["type"] == "loc_frame":
return descriptor
rcut = descriptor["rcut"]
tmp_sel = get_sel(
jdata, rcut, one_type=descriptor["type"] in ("se_atten",)
) # [38 72]，每个原子截断半径内，最多的邻域原子个数
tmp_sel = get_sel(jdata, rcut, one_type=descriptor["type"] in ("se_atten",))
sel = descriptor["sel"] # [46, 92]
if isinstance(sel, int):
# convert to list and finnally convert back to int
Expand Down
5 changes: 1 addition & 4 deletions deepmd/env.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -372,10 +372,7 @@ def get_module(module_name: str) -> "ModuleType":
raise FileNotFoundError(f"module {module_name} does not exist")
else:
try:
# module = tf.load_op_library(str(module_file))
import paddle_deepmd_lib

module = paddle_deepmd_lib
import paddle_deepmd_lib as module

except tf.errors.NotFoundError as e:
# check CXX11_ABI_FLAG is compatiblity
Expand Down
55 changes: 4 additions & 51 deletions deepmd/fit/ener.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -452,7 +452,6 @@ def _build_lower(
bias_atom_e=0.0,
type_suffix="",
suffix="",
# reuse=None,
type_i=None,
):
# cut-out inputs
Expand Down Expand Up @@ -483,47 +482,15 @@ def _build_lower(
ext_aparam = paddle.cast(ext_aparam, self.fitting_precision)
layer = paddle.concat([layer, ext_aparam], axis=1)

# if nvnmd_cfg.enable:
# one_layer = one_layer_nvnmd
# else:
# one_layer = one_layer_deepmd
for ii in range(0, len(self.n_neuron)):
# if self.layer_name is not None and self.layer_name[ii] is not None:
# layer_suffix = "share_" + self.layer_name[ii] + type_suffix
# layer_reuse = tf.AUTO_REUSE
# else:
# layer_suffix = "layer_" + str(ii) + type_suffix + suffix
# layer_reuse = reuse
if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii - 1]:
layer += self.one_layers[type_i][ii](layer)
else:
layer = self.one_layers[type_i][ii](layer)
# print(f"use {ii} of {len(self.one_layers)}_{type_i}")
# if (not self.uniform_seed) and (self.seed is not None):
# self.seed += self.seed_shift
# if self.layer_name is not None and self.layer_name[-1] is not None:
# layer_suffix = "share_" + self.layer_name[-1] + type_suffix
# layer_reuse = tf.AUTO_REUSE
# else:
# layer_suffix = "final_layer" + type_suffix + suffix
# layer_reuse = reuse
if (not self.uniform_seed) and (self.seed is not None):
self.seed += self.seed_shift
final_layer = self.final_layers[type_i](
layer,
# 1,
# activation_fn=None,
# bavg=bias_atom_e,
# name=layer_suffix,
# reuse=layer_reuse,
# seed=self.seed,
# precision=self.fitting_precision,
# trainable=self.trainable[-1],
# uniform_seed=self.uniform_seed,
# initial_variables=self.fitting_net_variables,
# mixed_prec=self.mixed_prec,
# final_layer=True,
)

final_layer = self.final_layers[type_i](layer)
if (not self.uniform_seed) and (self.seed is not None):
self.seed += self.seed_shift

Expand Down Expand Up @@ -646,18 +613,6 @@ def forward(
start_index = 0
outs_list = []
for type_i in range(ntypes_atom):
# final_layer = inputs
# for layer_j in range(type_i * ntypes_atom, (type_i + 1) * ntypes_atom):
# final_layer = self.one_layers[layer_j](final_layer)
# final_layer = self.final_layers[type_i](final_layer)
# print(final_layer.shape)

# # concat the results
# if type_i < len(self.atom_ener) and self.atom_ener[type_i] is not None:
# zero_layer = inputs_zero
# for layer_j in range(type_i * ntypes_atom, (type_i + 1) * ntypes_atom):
# zero_layer = self.one_layers[layer_j](zero_layer)
# zero_layer = self.final_layers[type_i](zero_layer)
final_layer = self._build_lower(
start_index,
natoms[2 + type_i],
Expand All @@ -667,7 +622,6 @@ def forward(
bias_atom_e=0.0,
type_suffix="_type_" + str(type_i),
suffix=suffix,
# reuse=reuse,
type_i=type_i,
)
# concat the results
Expand All @@ -681,13 +635,12 @@ def forward(
bias_atom_e=0.0,
type_suffix="_type_" + str(type_i),
suffix=suffix,
# reuse=True,
type_i=type_i,
)
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 @@ -734,7 +687,7 @@ def forward(
),
[paddle.shape(inputs)[0], paddle.sum(natoms[2 : 2 + ntypes_atom]).item()],
)
outs = outs + self.add_type # 类型编码(类似于transformer的位置编码，每种类型自己有一个特征，加到原特征上)
outs = outs + self.add_type
outs *= atype_filter
self.atom_ener_after = outs

Expand Down
13 changes: 6 additions & 7 deletions deepmd/infer/deep_eval.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -327,18 +327,18 @@ def __init__(
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}"
)
# 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}"
)
# 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
)
Expand Down Expand Up @@ -492,7 +492,6 @@ def _load_graph(
# # print(param.shape)
# if param.shape == (2,):
# print(constant_op.outputs[0], param)

return graph

@staticmethod
Expand Down
Loading

0 comments on commit 52638f2

Please sign in to comment.