Skip to content

GitLab

"examples/dreambooth/train_dreambooth_lora_lumina2.py" did not exist on "9408aa2dfc215c77ca40dd89fe4fc33f0d3826b5"
Commit 56d2b104 authored by zhangqha's avatar zhangqha
Browse files

BladeDISC DeePMD code update

parent 6b33aeb8
Pipeline #180 failed with stages
in 0 seconds
Hide whitespace changes
Inline Side-by-side
Showing with 0 additions and 2927 deletions
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/deep_tensor.py deleted 100644 → 0
View file @ 6b33aeb8
import os
from typing import List, Optional, TYPE_CHECKING, Tuple
import numpy as np
from deepmd.common import make_default_mesh
from deepmd.env import default_tf_session_config, tf
from deepmd.infer.deep_eval import DeepEval
from deepmd.utils.sess import run_sess
if TYPE_CHECKING:
from pathlib import Path
class DeepTensor(DeepEval):
"""Evaluates a tensor model.
Parameters
----------
model_file: str
The name of the frozen model file.
load_prefix: str
The prefix in the load computational graph
default_tf_graph : bool
If uses the default tf graph, otherwise build a new tf graph for evaluation
"""
tensors = {
# descriptor attrs
"t_ntypes": "descrpt_attr/ntypes:0",
"t_rcut": "descrpt_attr/rcut:0",
# model attrs
"t_tmap": "model_attr/tmap:0",
"t_sel_type": "model_attr/sel_type:0",
"t_ouput_dim": "model_attr/output_dim:0",
# inputs
"t_coord": "t_coord:0",
"t_type": "t_type:0",
"t_natoms": "t_natoms:0",
"t_box": "t_box:0",
"t_mesh": "t_mesh:0",
}
def __init__(
self,
model_file: "Path",
load_prefix: str = 'load',
default_tf_graph: bool = False
) -> None:
"""Constructor"""
DeepEval.__init__(
self,
model_file,
load_prefix=load_prefix,
default_tf_graph=default_tf_graph
)
# check model type
model_type = self.tensors["t_tensor"][2:-2]
assert self.model_type == model_type, \
f"expect {model_type} model but got {self.model_type}"
# now load tensors to object attributes
for attr_name, tensor_name in self.tensors.items():
self._get_tensor(tensor_name, attr_name)
# load optional tensors if possible
optional_tensors = {
"t_global_tensor": f"o_global_{model_type}:0",
"t_force": "o_force:0",
"t_virial": "o_virial:0",
"t_atom_virial": "o_atom_virial:0"
}
try:
# first make sure these tensor all exists (but do not modify self attr)
for attr_name, tensor_name in optional_tensors.items():
self._get_tensor(tensor_name)
# then put those into self.attrs
for attr_name, tensor_name in optional_tensors.items():
self._get_tensor(tensor_name, attr_name)
except KeyError:
self._support_gfv = False
else:
self.tensors.update(optional_tensors)
self._support_gfv = True
self._run_default_sess()
self.tmap = self.tmap.decode('UTF-8').split()
def _run_default_sess(self):
[self.ntypes, self.rcut, self.tmap, self.tselt, self.output_dim] \
= run_sess(self.sess,
[self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type, self.t_ouput_dim]
)
def get_ntypes(self) -> int:
"""Get the number of atom types of this model."""
return self.ntypes
def get_rcut(self) -> float:
"""Get the cut-off radius of this model."""
return self.rcut
def get_type_map(self) -> List[int]:
"""Get the type map (element name of the atom types) of this model."""
return self.tmap
def get_sel_type(self) -> List[int]:
"""Get the selected atom types of this model."""
return self.tselt
def get_dim_fparam(self) -> int:
"""Get the number (dimension) of frame parameters of this DP."""
return self.dfparam
def get_dim_aparam(self) -> int:
"""Get the number (dimension) of atomic parameters of this DP."""
return self.daparam
def eval(
self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
atomic: bool = True,
fparam: Optional[np.ndarray] = None,
aparam: Optional[np.ndarray] = None,
efield: Optional[np.ndarray] = None
) -> np.ndarray:
"""Evaluate the model.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
atomic
If True (default), return the atomic tensor
Otherwise return the global tensor
fparam
Not used in this model
aparam
Not used in this model
efield
Not used in this model
Returns
-------
tensor
The returned tensor
If atomic == False then of size nframes x output_dim
else of size nframes x natoms x output_dim
"""
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype = int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
# sort inputs
coords, atom_types, imap, sel_at, sel_imap = self.sort_input(coords, atom_types, sel_atoms = self.get_sel_type())
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert(natoms_vec[0] == natoms)
# evaluate
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type ] = np.tile(atom_types, [nframes,1]).reshape([-1])
feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
feed_dict_test[self.t_box ] = np.reshape(cells , [-1])
if pbc:
feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
if atomic:
assert "global" not in self.model_type, \
f"cannot do atomic evaluation with model type {self.model_type}"
t_out = [self.t_tensor]
else:
assert self._support_gfv or "global" in self.model_type, \
f"do not support global tensor evaluation with old {self.model_type} model"
t_out = [self.t_global_tensor if self._support_gfv else self.t_tensor]
v_out = self.sess.run (t_out, feed_dict = feed_dict_test)
tensor = v_out[0]
# reverse map of the outputs
if atomic:
tensor = np.array(tensor)
tensor = self.reverse_map(np.reshape(tensor, [nframes,-1,self.output_dim]), sel_imap)
tensor = np.reshape(tensor, [nframes, len(sel_at), self.output_dim])
else:
tensor = np.reshape(tensor, [nframes, self.output_dim])
return tensor
def eval_full(
self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
atomic: bool = False,
fparam: Optional[np.array] = None,
aparam: Optional[np.array] = None,
efield: Optional[np.array] = None
) -> Tuple[np.ndarray, ...]:
"""Evaluate the model with interface similar to the energy model.
Will return global tensor, component-wise force and virial
and optionally atomic tensor and atomic virial.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
atomic
Whether to calculate atomic tensor and virial
fparam
Not used in this model
aparam
Not used in this model
efield
Not used in this model
Returns
-------
tensor
The global tensor.
shape: [nframes x nout]
force
The component-wise force (negative derivative) on each atom.
shape: [nframes x nout x natoms x 3]
virial
The component-wise virial of the tensor.
shape: [nframes x nout x 9]
atom_tensor
The atomic tensor. Only returned when atomic == True
shape: [nframes x natoms x nout]
atom_virial
The atomic virial. Only returned when atomic == True
shape: [nframes x nout x natoms x 9]
"""
assert self._support_gfv, \
f"do not support eval_full with old tensor model"
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype = int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
nout = self.output_dim
# sort inputs
coords, atom_types, imap, sel_at, sel_imap = self.sort_input(coords, atom_types, sel_atoms = self.get_sel_type())
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert(natoms_vec[0] == natoms)
# evaluate
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type ] = np.tile(atom_types, [nframes,1]).reshape([-1])
feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
feed_dict_test[self.t_box ] = np.reshape(cells , [-1])
if pbc:
feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
t_out = [self.t_global_tensor,
self.t_force,
self.t_virial]
if atomic :
t_out += [self.t_tensor,
self.t_atom_virial]
v_out = self.sess.run (t_out, feed_dict = feed_dict_test)
gt = v_out[0] # global tensor
force = v_out[1]
virial = v_out[2]
if atomic:
at = v_out[3] # atom tensor
av = v_out[4] # atom virial
# please note here the shape are wrong!
force = self.reverse_map(np.reshape(force, [nframes*nout, natoms ,3]), imap)
if atomic:
at = self.reverse_map(np.reshape(at, [nframes, len(sel_at), nout]), sel_imap)
av = self.reverse_map(np.reshape(av, [nframes*nout, natoms, 9]), imap)
# make sure the shapes are correct here
gt = np.reshape(gt, [nframes, nout])
force = np.reshape(force, [nframes, nout, natoms, 3])
virial = np.reshape(virial, [nframes, nout, 9])
if atomic:
at = np.reshape(at, [nframes, len(sel_at), self.output_dim])
av = np.reshape(av, [nframes, nout, natoms, 9])
return gt, force, virial, at, av
else:
return gt, force, virial
\ No newline at end of file
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/deep_wfc.py deleted 100644 → 0
View file @ 6b33aeb8
from deepmd.infer.deep_tensor import DeepTensor
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from pathlib import Path
class DeepWFC(DeepTensor):
"""Constructor.
Parameters
----------
model_file : Path
The name of the frozen model file.
load_prefix: str
The prefix in the load computational graph
default_tf_graph : bool
If uses the default tf graph, otherwise build a new tf graph for evaluation
Warnings
--------
For developers: `DeepTensor` initializer must be called at the end after
`self.tensors` are modified because it uses the data in `self.tensors` dict.
Do not chanage the order!
"""
def __init__(
self, model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False
) -> None:
# use this in favor of dict update to move attribute from class to
# instance namespace
self.tensors = dict(
{
# output tensor
"t_tensor": "o_wfc:0",
},
**self.tensors
)
DeepTensor.__init__(
self,
model_file,
load_prefix=load_prefix,
default_tf_graph=default_tf_graph,
)
def get_dim_fparam(self) -> int:
"""Unsupported in this model."""
raise NotImplementedError("This model type does not support this attribute")
def get_dim_aparam(self) -> int:
"""Unsupported in this model."""
raise NotImplementedError("This model type does not support this attribute")
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/ewald_recp.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.common import ClassArg
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import global_cvt_2_ener_float
from deepmd.env import op_module
from deepmd.env import default_tf_session_config
from deepmd.utils.sess import run_sess
class EwaldRecp () :
"""
Evaluate the reciprocal part of the Ewald sum
"""
def __init__(self,
hh,
beta):
"""
Constructor
Parameters
----------
hh
Grid spacing of the reciprocal part of Ewald sum. Unit: A
beta
Splitting parameter of the Ewald sum. Unit: A^{-1}
"""
self.hh = hh
self.beta = beta
with tf.Graph().as_default() as graph:
# place holders
self.t_nloc = tf.placeholder(tf.int32, [1], name = "t_nloc")
self.t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_coord')
self.t_charge = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_charge')
self.t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_box')
# output
self.t_energy, self.t_force, self.t_virial \
= op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box,
ewald_h = self.hh,
ewald_beta = self.beta)
self.sess = tf.Session(graph=graph, config=default_tf_session_config)
def eval(self,
coord : np.ndarray,
charge : np.ndarray,
box : np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray] :
"""
Evaluate
Parameters
----------
coord
The coordinates of atoms
charge
The atomic charge
box
The simulation region. PBC is assumed
Returns
-------
e
The energy
f
The force
v
The virial
"""
coord = np.array(coord)
charge = np.array(charge)
box = np.array(box)
nframes = charge.shape[0]
natoms = charge.shape[1]
coord = np.reshape(coord, [nframes * 3 * natoms])
charge = np.reshape(charge, [nframes * natoms])
box = np.reshape(box, [nframes * 9])
[energy, force, virial] \
= run_sess(self.sess, [self.t_energy, self.t_force, self.t_virial],
feed_dict = {
self.t_coord: coord,
self.t_charge: charge,
self.t_box: box,
self.t_nloc: [natoms],
})
return energy, force, virial
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/model_devi.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from .deep_pot import DeepPot
from ..utils.data import DeepmdData
from ..utils.batch_size import AutoBatchSize
from deepmd.common import expand_sys_str
def calc_model_devi_f(fs: np.ndarray):
'''
Parameters
----------
fs : numpy.ndarray
size of `n_models x n_frames x n_atoms x 3`
'''
fs_devi = np.linalg.norm(np.std(fs, axis=0), axis=-1)
max_devi_f = np.max(fs_devi, axis=-1)
min_devi_f = np.min(fs_devi, axis=-1)
avg_devi_f = np.mean(fs_devi, axis=-1)
return max_devi_f, min_devi_f, avg_devi_f
def calc_model_devi_e(es: np.ndarray):
'''
Parameters
----------
es : numpy.ndarray
size of `n_models x n_frames x n_atoms
'''
es_devi = np.std(es, axis=0)
max_devi_e = np.max(es_devi, axis=1)
min_devi_e = np.min(es_devi, axis=1)
avg_devi_e = np.mean(es_devi, axis=1)
return max_devi_e, min_devi_e, avg_devi_e
def calc_model_devi_v(vs: np.ndarray):
'''
Parameters
----------
vs : numpy.ndarray
size of `n_models x n_frames x 9`
'''
vs_devi = np.std(vs, axis=0)
max_devi_v = np.max(vs_devi, axis=-1)
min_devi_v = np.min(vs_devi, axis=-1)
avg_devi_v = np.linalg.norm(vs_devi, axis=-1) / 3
return max_devi_v, min_devi_v, avg_devi_v
def write_model_devi_out(devi: np.ndarray, fname: str, header: str=""):
'''
Parameters
----------
devi : numpy.ndarray
the first column is the steps index
fname : str
the file name to dump
header : str, default=""
the header to dump
'''
assert devi.shape[1] == 7
header = "%s\n%10s" % (header, "step")
for item in 'vf':
header += "%19s%19s%19s" % (f"max_devi_{item}", f"min_devi_{item}", f"avg_devi_{item}")
with open(fname, "ab") as fp:
np.savetxt(fp,
devi,
fmt=['%12d'] + ['%19.6e' for _ in range(6)],
delimiter='',
header=header)
return devi
def _check_tmaps(tmaps, ref_tmap=None):
'''
Check whether type maps are identical
'''
assert isinstance(tmaps, list)
if ref_tmap is None:
ref_tmap = tmaps[0]
assert isinstance(ref_tmap, list)
flag = True
for tmap in tmaps:
if tmap != ref_tmap:
flag = False
break
return flag
def calc_model_devi(coord,
box,
atype,
models,
fname=None,
frequency=1,
):
'''
Python interface to calculate model deviation
Parameters
-----------
coord : numpy.ndarray, `n_frames x n_atoms x 3`
Coordinates of system to calculate
box : numpy.ndarray or None, `n_frames x 3 x 3`
Box to specify periodic boundary condition. If None, no pbc will be used
atype : numpy.ndarray, `n_atoms x 1`
Atom types
models : list of DeepPot models
Models used to evaluate deviation
fname : str or None
File to dump results, default None
frequency : int
Steps between frames (if the system is given by molecular dynamics engine), default 1
Returns
-------
model_devi : numpy.ndarray, `n_frames x 7`
Model deviation results. The first column is index of steps, the other 6 columns are
max_devi_v, min_devi_v, avg_devi_v, max_devi_f, min_devi_f, avg_devi_f.
Examples
--------
>>> from deepmd.infer import calc_model_devi
>>> from deepmd.infer import DeepPot as DP
>>> import numpy as np
>>> coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
>>> cell = np.diag(10 * np.ones(3)).reshape([1, -1])
>>> atype = [1,0,1]
>>> graphs = [DP("graph.000.pb"), DP("graph.001.pb")]
>>> model_devi = calc_model_devi(coord, cell, atype, graphs)
'''
if box is not None:
nopbc = True
else:
nopbc = False
forces = []
virials = []
for dp in models:
ret = dp.eval(
coord,
box,
atype,
)
forces.append(ret[1])
virials.append(ret[2] / len(atype))
forces = np.array(forces)
virials = np.array(virials)
devi = [np.arange(coord.shape[0]) * frequency]
devi += list(calc_model_devi_v(virials))
devi += list(calc_model_devi_f(forces))
devi = np.vstack(devi).T
if fname:
write_model_devi_out(devi, fname)
return devi
def make_model_devi(
*,
models: list,
system: str,
set_prefix: str,
output: str,
frequency: int,
**kwargs
):
'''
Make model deviation calculation
Parameters
----------
models: list
A list of paths of models to use for making model deviation
system: str
The path of system to make model deviation calculation
set_prefix: str
The set prefix of the system
output: str
The output file for model deviation results
frequency: int
The number of steps that elapse between writing coordinates
in a trajectory by a MD engine (such as Gromacs / Lammps).
This paramter is used to determine the index in the output file.
'''
auto_batch_size = AutoBatchSize()
# init models
dp_models = [DeepPot(model, auto_batch_size=auto_batch_size) for model in models]
# check type maps
tmaps = [dp.get_type_map() for dp in dp_models]
if _check_tmaps(tmaps):
tmap = tmaps[0]
else:
raise RuntimeError("The models does not have the same type map.")
all_sys = expand_sys_str(system)
if len(all_sys) == 0:
raise RuntimeError("Did not find valid system")
devis_coll = []
for system in all_sys:
# create data-system
dp_data = DeepmdData(system, set_prefix, shuffle_test=False, type_map=tmap)
data_sets = [dp_data._load_set(set_name) for set_name in dp_data.dirs]
nframes_tot = 0
devis = []
for data in data_sets:
coord = data["coord"]
box = data["box"]
atype = data["type"][0]
if not dp_data.pbc:
box = None
devi = calc_model_devi(coord, box, atype, dp_models)
nframes_tot += coord.shape[0]
devis.append(devi)
devis = np.vstack(devis)
devis[:, 0] = np.arange(nframes_tot) * frequency
write_model_devi_out(devis, output, header=system)
devis_coll.append(devis)
return devis_coll
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/loggers/__init__.py deleted 100644 → 0
View file @ 6b33aeb8
"""Module taking care of logging duties."""
from .loggers import set_log_handles
__all__ = ["set_log_handles"]
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/loggers/loggers.py deleted 100644 → 0
View file @ 6b33aeb8
"""Logger initialization for package."""
import logging
import os
from typing import TYPE_CHECKING, Optional
if TYPE_CHECKING:
from pathlib import Path
from mpi4py import MPI
_MPI_APPEND_MODE = MPI.MODE_CREATE | MPI.MODE_APPEND
logging.getLogger(__name__)
__all__ = ["set_log_handles"]
# logger formater
FFORMATTER = logging.Formatter(
"[%(asctime)s] %(app_name)s %(levelname)-7s %(name)-45s %(message)s"
)
CFORMATTER = logging.Formatter(
# "%(app_name)s %(levelname)-7s |-> %(name)-45s %(message)s"
"%(app_name)s %(levelname)-7s %(message)s"
)
FFORMATTER_MPI = logging.Formatter(
"[%(asctime)s] %(app_name)s rank:%(rank)-2s %(levelname)-7s %(name)-45s %(message)s"
)
CFORMATTER_MPI = logging.Formatter(
# "%(app_name)s rank:%(rank)-2s %(levelname)-7s |-> %(name)-45s %(message)s"
"%(app_name)s rank:%(rank)-2s %(levelname)-7s %(message)s"
)
class _AppFilter(logging.Filter):
"""Add field `app_name` to log messages."""
def filter(self, record):
record.app_name = "DEEPMD"
return True
class _MPIRankFilter(logging.Filter):
"""Add MPI rank number to log messages, adds field `rank`."""
def __init__(self, rank: int) -> None:
super().__init__(name="MPI_rank_id")
self.mpi_rank = str(rank)
def filter(self, record):
record.rank = self.mpi_rank
return True
class _MPIMasterFilter(logging.Filter):
"""Filter that lets through only messages emited from rank==0."""
def __init__(self, rank: int) -> None:
super().__init__(name="MPI_master_log")
self.mpi_rank = rank
def filter(self, record):
if self.mpi_rank == 0:
return True
else:
return False
class _MPIFileStream:
"""Wrap MPI.File` so it has the same API as python file streams.
Parameters
----------
filename : Path
disk location of the file stream
MPI : MPI
MPI communicator object
mode : str, optional
file write mode, by default _MPI_APPEND_MODE
"""
def __init__(
self, filename: "Path", MPI: "MPI", mode: str = "_MPI_APPEND_MODE"
) -> None:
self.stream = MPI.File.Open(MPI.COMM_WORLD, filename, mode)
self.stream.Set_atomicity(True)
self.name = "MPIfilestream"
def write(self, msg: str):
"""Write to MPI shared file stream.
Parameters
----------
msg : str
message to write
"""
b = bytearray()
b.extend(map(ord, msg))
self.stream.Write_shared(b)
def close(self):
"""Synchronize and close MPI file stream."""
self.stream.Sync()
self.stream.Close()
class _MPIHandler(logging.FileHandler):
"""Emulate `logging.FileHandler` with MPI shared File that all ranks can write to.
Parameters
----------
filename : Path
file path
MPI : MPI
MPI communicator object
mode : str, optional
file access mode, by default "_MPI_APPEND_MODE"
"""
def __init__(
self,
filename: "Path",
MPI: "MPI",
mode: str = "_MPI_APPEND_MODE",
) -> None:
self.MPI = MPI
super().__init__(filename, mode=mode, encoding=None, delay=False)
def _open(self):
return _MPIFileStream(self.baseFilename, self.MPI, self.mode)
def setStream(self, stream):
"""Stream canot be reasigned in MPI mode."""
raise NotImplementedError("Unable to do for MPI file handler!")
def set_log_handles(
level: int,
log_path: Optional["Path"] = None,
mpi_log: Optional[str] = None
):
"""Set desired level for package loggers and add file handlers.
Parameters
----------
level: int
logging level
log_path: Optional[str]
path to log file, if None logs will be send only to console. If the parent
directory does not exist it will be automatically created, by default None
mpi_log : Optional[str], optional
mpi log type. Has three options. `master` will output logs to file and console
only from rank==0. `collect` will write messages from all ranks to one file
opened under rank==0 and to console. `workers` will open one log file for each
worker designated by its rank, console behaviour is the same as for `collect`.
If this argument is specified, package 'mpi4py' must be already installed.
by default None
Raises
------
RuntimeError
If the argument `mpi_log` is specified, package `mpi4py` is not installed.
References
----------
https://groups.google.com/g/mpi4py/c/SaNzc8bdj6U
https://stackoverflow.com/questions/35869137/avoid-tensorflow-print-on-standard-error
https://stackoverflow.com/questions/56085015/suppress-openmp-debug-messages-when-running-tensorflow-on-cpu
Notes
-----
Logging levels:
+---------+--------------+----------------+----------------+----------------+
| | our notation | python logging | tensorflow cpp | OpenMP |
+=========+==============+================+================+================+
| debug | 10 | 10 | 0 | 1/on/true/yes |
+---------+--------------+----------------+----------------+----------------+
| info | 20 | 20 | 1 | 0/off/false/no |
+---------+--------------+----------------+----------------+----------------+
| warning | 30 | 30 | 2 | 0/off/false/no |
+---------+--------------+----------------+----------------+----------------+
| error | 40 | 40 | 3 | 0/off/false/no |
+---------+--------------+----------------+----------------+----------------+
"""
# silence logging for OpenMP when running on CPU if level is any other than debug
if level <= 10:
os.environ["KMP_WARNINGS"] = "FALSE"
# set TF cpp internal logging level
os.environ['TF_CPP_MIN_LOG_LEVEL'] = str(int((level / 10) - 1))
# get root logger
root_log = logging.getLogger()
# remove all old handlers
root_log.setLevel(level)
for hdlr in root_log.handlers[:]:
root_log.removeHandler(hdlr)
# check if arguments are present
MPI = None
if mpi_log:
try:
from mpi4py import MPI
except ImportError as e:
raise RuntimeError("You cannot specify 'mpi_log' when mpi4py not installed") from e
# * add console handler ************************************************************
ch = logging.StreamHandler()
if MPI:
rank = MPI.COMM_WORLD.Get_rank()
if mpi_log == "master":
ch.setFormatter(CFORMATTER)
ch.addFilter(_MPIMasterFilter(rank))
else:
ch.setFormatter(CFORMATTER_MPI)
ch.addFilter(_MPIRankFilter(rank))
else:
ch.setFormatter(CFORMATTER)
ch.setLevel(level)
ch.addFilter(_AppFilter())
root_log.addHandler(ch)
# * add file handler ***************************************************************
if log_path:
# create directory
log_path.parent.mkdir(exist_ok=True, parents=True)
fh = None
if mpi_log == "master":
rank = MPI.COMM_WORLD.Get_rank()
if rank == 0:
fh = logging.FileHandler(log_path, mode="w")
fh.addFilter(_MPIMasterFilter(rank))
fh.setFormatter(FFORMATTER)
elif mpi_log == "collect":
rank = MPI.COMM_WORLD.Get_rank()
fh = _MPIHandler(log_path, MPI, mode=MPI.MODE_WRONLY | MPI.MODE_CREATE)
fh.addFilter(_MPIRankFilter(rank))
fh.setFormatter(FFORMATTER_MPI)
elif mpi_log == "workers":
rank = MPI.COMM_WORLD.Get_rank()
# if file has suffix than inser rank number before suffix
# e.g deepmd.log -> deepmd_<rank>.log
# if no suffix is present, insert rank as suffix
# e.g. deepmdlog -> deepmdlog.<rank>
if log_path.suffix:
worker_log = (log_path.parent / f"{log_path.stem}_{rank}").with_suffix(
log_path.suffix
)
else:
worker_log = log_path.with_suffix(f".{rank}")
fh = logging.FileHandler(worker_log, mode="w")
fh.setFormatter(FFORMATTER)
else:
fh = logging.FileHandler(log_path, mode="w")
fh.setFormatter(FFORMATTER)
if fh:
fh.setLevel(level)
fh.addFilter(_AppFilter())
root_log.addHandler(fh)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/loss/__init__.py deleted 100644 → 0
View file @ 6b33aeb8
from .ener import EnerStdLoss
from .ener import EnerDipoleLoss
from .tensor import TensorLoss
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/loss/ener.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from deepmd.env import tf
from deepmd.common import ClassArg, add_data_requirement
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import global_cvt_2_ener_float
from deepmd.utils.sess import run_sess
from .loss import Loss
class EnerStdLoss (Loss) :
"""
Standard loss function for DP models
Parameters
----------
enable_atom_ener_coeff : bool
if true, the energy will be computed as \sum_i c_i E_i
"""
def __init__ (self,
starter_learning_rate : float,
start_pref_e : float = 0.02,
limit_pref_e : float = 1.00,
start_pref_f : float = 1000,
limit_pref_f : float = 1.00,
start_pref_v : float = 0.0,
limit_pref_v : float = 0.0,
start_pref_ae : float = 0.0,
limit_pref_ae : float = 0.0,
start_pref_pf : float = 0.0,
limit_pref_pf : float = 0.0,
relative_f : float = None,
enable_atom_ener_coeff: bool=False,
) -> None:
self.starter_learning_rate = starter_learning_rate
self.start_pref_e = start_pref_e
self.limit_pref_e = limit_pref_e
self.start_pref_f = start_pref_f
self.limit_pref_f = limit_pref_f
self.start_pref_v = start_pref_v
self.limit_pref_v = limit_pref_v
self.start_pref_ae = start_pref_ae
self.limit_pref_ae = limit_pref_ae
self.start_pref_pf = start_pref_pf
self.limit_pref_pf = limit_pref_pf
self.relative_f = relative_f
self.enable_atom_ener_coeff = enable_atom_ener_coeff
self.has_e = (self.start_pref_e != 0.0 or self.limit_pref_e != 0.0)
self.has_f = (self.start_pref_f != 0.0 or self.limit_pref_f != 0.0)
self.has_v = (self.start_pref_v != 0.0 or self.limit_pref_v != 0.0)
self.has_ae = (self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0)
self.has_pf = (self.start_pref_pf != 0.0 or self.limit_pref_pf != 0.0)
# data required
add_data_requirement('energy', 1, atomic=False, must=False, high_prec=True)
add_data_requirement('force', 3, atomic=True, must=False, high_prec=False)
add_data_requirement('virial', 9, atomic=False, must=False, high_prec=False)
add_data_requirement('atom_ener', 1, atomic=True, must=False, high_prec=False)
add_data_requirement('atom_pref', 1, atomic=True, must=False, high_prec=False, repeat=3)
if self.enable_atom_ener_coeff:
add_data_requirement('atom_ener_coeff', 1, atomic=True, must=False, high_prec=False, default=1.)
def build (self,
learning_rate,
natoms,
model_dict,
label_dict,
suffix):
energy = model_dict['energy']
force = model_dict['force']
virial = model_dict['virial']
atom_ener = model_dict['atom_ener']
energy_hat = label_dict['energy']
force_hat = label_dict['force']
virial_hat = label_dict['virial']
atom_ener_hat = label_dict['atom_ener']
atom_pref = label_dict['atom_pref']
find_energy = label_dict['find_energy']
find_force = label_dict['find_force']
find_virial = label_dict['find_virial']
find_atom_ener = label_dict['find_atom_ener']
find_atom_pref = label_dict['find_atom_pref']
if self.enable_atom_ener_coeff:
# when ener_coeff (\nu) is defined, the energy is defined as
# E = \sum_i \nu_i E_i
# instead of the sum of atomic energies.
#
# A case is that we want to train reaction energy
# A + B -> C + D
# E = - E(A) - E(B) + E(C) + E(D)
# A, B, C, D could be put far away from each other
atom_ener_coeff = label_dict['atom_ener_coeff']
atom_ener_coeff = tf.reshape(atom_ener_coeff, tf.shape(atom_ener))
energy = tf.reduce_sum(atom_ener_coeff * atom_ener, 1)
l2_ener_loss = tf.reduce_mean( tf.square(energy - energy_hat), name='l2_'+suffix)
force_reshape = tf.reshape (force, [-1])
force_hat_reshape = tf.reshape (force_hat, [-1])
atom_pref_reshape = tf.reshape (atom_pref, [-1])
diff_f = force_hat_reshape - force_reshape
if self.relative_f is not None:
force_hat_3 = tf.reshape(force_hat, [-1, 3])
norm_f = tf.reshape(tf.norm(force_hat_3, axis = 1), [-1, 1]) + self.relative_f
diff_f_3 = tf.reshape(diff_f, [-1, 3])
diff_f_3 = diff_f_3 / norm_f
diff_f = tf.reshape(diff_f_3, [-1])
l2_force_loss = tf.reduce_mean(tf.square(diff_f), name = "l2_force_" + suffix)
l2_pref_force_loss = tf.reduce_mean(tf.multiply(tf.square(diff_f), atom_pref_reshape), name = "l2_pref_force_" + suffix)
virial_reshape = tf.reshape (virial, [-1])
virial_hat_reshape = tf.reshape (virial_hat, [-1])
l2_virial_loss = tf.reduce_mean (tf.square(virial_hat_reshape - virial_reshape), name = "l2_virial_" + suffix)
atom_ener_reshape = tf.reshape (atom_ener, [-1])
atom_ener_hat_reshape = tf.reshape (atom_ener_hat, [-1])
l2_atom_ener_loss = tf.reduce_mean (tf.square(atom_ener_hat_reshape - atom_ener_reshape), name = "l2_atom_ener_" + suffix)
atom_norm = 1./ global_cvt_2_tf_float(natoms[0])
atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0])
pref_e = global_cvt_2_ener_float(find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate) )
pref_f = global_cvt_2_tf_float(find_force * (self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * learning_rate / self.starter_learning_rate) )
pref_v = global_cvt_2_tf_float(find_virial * (self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * learning_rate / self.starter_learning_rate) )
pref_ae= global_cvt_2_tf_float(find_atom_ener * (self.limit_pref_ae+ (self.start_pref_ae-self.limit_pref_ae) * learning_rate / self.starter_learning_rate) )
pref_pf= global_cvt_2_tf_float(find_atom_pref * (self.limit_pref_pf+ (self.start_pref_pf-self.limit_pref_pf) * learning_rate / self.starter_learning_rate) )
l2_loss = 0
more_loss = {}
if self.has_e :
l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
more_loss['l2_ener_loss'] = l2_ener_loss
if self.has_f :
l2_loss += global_cvt_2_ener_float(pref_f * l2_force_loss)
more_loss['l2_force_loss'] = l2_force_loss
if self.has_v :
l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
more_loss['l2_virial_loss'] = l2_virial_loss
if self.has_ae :
l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss)
more_loss['l2_atom_ener_loss'] = l2_atom_ener_loss
if self.has_pf :
l2_loss += global_cvt_2_ener_float(pref_pf * l2_pref_force_loss)
more_loss['l2_pref_force_loss'] = l2_pref_force_loss
# only used when tensorboard was set as true
self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', global_cvt_2_tf_float(tf.sqrt(l2_ener_loss)) / global_cvt_2_tf_float(natoms[0]))
self.l2_loss_force_summary = tf.summary.scalar('l2_force_loss', tf.sqrt(l2_force_loss))
self.l2_loss_virial_summary = tf.summary.scalar('l2_virial_loss', tf.sqrt(l2_virial_loss) / global_cvt_2_tf_float(natoms[0]))
self.l2_l = l2_loss
self.l2_more = more_loss
return l2_loss, more_loss
def eval(self, sess, feed_dict, natoms):
placeholder = self.l2_l
run_data = [
self.l2_l,
self.l2_more['l2_ener_loss'] if self.has_e else placeholder,
self.l2_more['l2_force_loss'] if self.has_f else placeholder,
self.l2_more['l2_virial_loss'] if self.has_v else placeholder,
self.l2_more['l2_atom_ener_loss'] if self.has_ae else placeholder,
self.l2_more['l2_pref_force_loss'] if self.has_pf else placeholder,
]
error, error_e, error_f, error_v, error_ae, error_pf = run_sess(sess, run_data, feed_dict=feed_dict)
results = {"natoms": natoms[0], "rmse": np.sqrt(error)}
if self.has_e:
results["rmse_e"] = np.sqrt(error_e) / natoms[0]
if self.has_ae:
results["rmse_ae"] = np.sqrt(error_ae)
if self.has_f:
results["rmse_f"] = np.sqrt(error_f)
if self.has_v:
results["rmse_v"] = np.sqrt(error_v) / natoms[0]
if self.has_pf:
results["rmse_pf"] = np.sqrt(error_pf)
return results
def print_header(self): # depreciated
prop_fmt = ' %11s %11s'
print_str = ''
print_str += prop_fmt % ('rmse_tst', 'rmse_trn')
if self.has_e :
print_str += prop_fmt % ('rmse_e_tst', 'rmse_e_trn')
if self.has_ae :
print_str += prop_fmt % ('rmse_ae_tst', 'rmse_ae_trn')
if self.has_f :
print_str += prop_fmt % ('rmse_f_tst', 'rmse_f_trn')
if self.has_v :
print_str += prop_fmt % ('rmse_v_tst', 'rmse_v_trn')
if self.has_pf :
print_str += prop_fmt % ('rmse_pf_tst', 'rmse_pf_trn')
return print_str
def print_on_training(self,
tb_writer,
cur_batch,
sess,
natoms,
feed_dict_test,
feed_dict_batch): # depreciated
run_data = [
self.l2_l,
self.l2_more['l2_ener_loss'],
self.l2_more['l2_force_loss'],
self.l2_more['l2_virial_loss'],
self.l2_more['l2_atom_ener_loss'],
self.l2_more['l2_pref_force_loss']
]
# first train data
train_out = run_sess(sess, run_data, feed_dict=feed_dict_batch)
error_train, error_e_train, error_f_train, error_v_train, error_ae_train, error_pf_train = train_out
# than test data, if tensorboard log writter is present, commpute summary
# and write tensorboard logs
if tb_writer:
summary_merged_op = tf.summary.merge([self.l2_loss_summary, self.l2_loss_ener_summary, self.l2_loss_force_summary, self.l2_loss_virial_summary])
run_data.insert(0, summary_merged_op)
test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
if tb_writer:
summary = test_out.pop(0)
tb_writer.add_summary(summary, cur_batch)
error_test, error_e_test, error_f_test, error_v_test, error_ae_test, error_pf_test = test_out
print_str = ""
prop_fmt = " %11.2e %11.2e"
print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
if self.has_e :
print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0])
if self.has_ae :
print_str += prop_fmt % (np.sqrt(error_ae_test), np.sqrt(error_ae_train))
if self.has_f :
print_str += prop_fmt % (np.sqrt(error_f_test), np.sqrt(error_f_train))
if self.has_v :
print_str += prop_fmt % (np.sqrt(error_v_test) / natoms[0], np.sqrt(error_v_train) / natoms[0])
if self.has_pf:
print_str += prop_fmt % (np.sqrt(error_pf_test), np.sqrt(error_pf_train))
return print_str
class EnerDipoleLoss (Loss) :
def __init__ (self,
starter_learning_rate : float,
start_pref_e : float = 0.1,
limit_pref_e : float = 1.0,
start_pref_ed : float = 1.0,
limit_pref_ed : float = 1.0
) -> None :
self.starter_learning_rate = kwarg['starter_learning_rate']
args = ClassArg()\
.add('start_pref_e', float, must = True, default = 0.1) \
.add('limit_pref_e', float, must = True, default = 1.00)\
.add('start_pref_ed', float, must = True, default = 1.00)\
.add('limit_pref_ed', float, must = True, default = 1.00)
class_data = args.parse(jdata)
self.start_pref_e = class_data['start_pref_e']
self.limit_pref_e = class_data['limit_pref_e']
self.start_pref_ed = class_data['start_pref_ed']
self.limit_pref_ed = class_data['limit_pref_ed']
# data required
add_data_requirement('energy', 1, atomic=False, must=True, high_prec=True)
add_data_requirement('energy_dipole', 3, atomic=False, must=True, high_prec=False)
def build (self,
learning_rate,
natoms,
model_dict,
label_dict,
suffix):
coord = model_dict['coord']
energy = model_dict['energy']
atom_ener = model_dict['atom_ener']
nframes = tf.shape(atom_ener)[0]
natoms = tf.shape(atom_ener)[1]
# build energy dipole
atom_ener0 = atom_ener - tf.reshape(tf.tile(tf.reshape(energy/global_cvt_2_ener_float(natoms), [-1, 1]), [1, natoms]), [nframes, natoms])
coord = tf.reshape(coord, [nframes, natoms, 3])
atom_ener0 = tf.reshape(atom_ener0, [nframes, 1, natoms])
ener_dipole = tf.matmul(atom_ener0, coord)
ener_dipole = tf.reshape(ener_dipole, [nframes, 3])
energy_hat = label_dict['energy']
ener_dipole_hat = label_dict['energy_dipole']
find_energy = label_dict['find_energy']
find_ener_dipole = label_dict['find_energy_dipole']
l2_ener_loss = tf.reduce_mean( tf.square(energy - energy_hat), name='l2_'+suffix)
ener_dipole_reshape = tf.reshape(ener_dipole, [-1])
ener_dipole_hat_reshape = tf.reshape(ener_dipole_hat, [-1])
l2_ener_dipole_loss = tf.reduce_mean( tf.square(ener_dipole_reshape - ener_dipole_hat_reshape), name='l2_'+suffix)
# atom_norm_ener = 1./ global_cvt_2_ener_float(natoms[0])
atom_norm_ener = 1./ global_cvt_2_ener_float(natoms)
pref_e = global_cvt_2_ener_float(find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate) )
pref_ed = global_cvt_2_tf_float(find_ener_dipole * (self.limit_pref_ed + (self.start_pref_ed - self.limit_pref_ed) * learning_rate / self.starter_learning_rate) )
l2_loss = 0
more_loss = {}
l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss)
more_loss['l2_ener_loss'] = l2_ener_loss
more_loss['l2_ener_dipole_loss'] = l2_ener_dipole_loss
self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', tf.sqrt(l2_ener_loss) / global_cvt_2_tf_float(natoms[0]))
self.l2_ener_dipole_loss_summary = tf.summary.scalar('l2_ener_dipole_loss', tf.sqrt(l2_ener_dipole_loss))
self.l2_l = l2_loss
self.l2_more = more_loss
return l2_loss, more_loss
def eval(self, sess, feed_dict, natoms):
run_data = [
self.l2_l,
self.l2_more['l2_ener_loss'],
self.l2_more['l2_ener_dipole_loss']
]
error, error_e, error_ed = run_sess(sess, run_data, feed_dict=feed_dict)
results = {
'natoms': natoms[0],
'rmse': np.sqrt(error),
'rmse_e': np.sqrt(error_e) / natoms[0],
'rmse_ed': np.sqrt(error_ed)
}
return results
@staticmethod
def print_header() : # depreciated
prop_fmt = ' %9s %9s'
print_str = ''
print_str += prop_fmt % ('l2_tst', 'l2_trn')
print_str += prop_fmt % ('l2_e_tst', 'l2_e_trn')
print_str += prop_fmt % ('l2_ed_tst', 'l2_ed_trn')
return print_str
def print_on_training(self,
tb_writer,
cur_batch,
sess,
natoms,
feed_dict_test,
feed_dict_batch): # depreciated
run_data = [
self.l2_l,
self.l2_more['l2_ener_loss'],
self.l2_more['l2_ener_dipole_loss']
]
# first train data
train_out = run_sess(sess, run_data, feed_dict=feed_dict_batch)
error_train, error_e_train, error_ed_train = train_out
# than test data, if tensorboard log writter is present, commpute summary
# and write tensorboard logs
if tb_writer:
summary_merged_op = tf.summary.merge([
self.l2_loss_summary,
self.l2_loss_ener_summary,
self.l2_ener_dipole_loss_summary
])
run_data.insert(0, summary_merged_op)
test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
if tb_writer:
summary = test_out.pop(0)
tb_writer.add_summary(summary, cur_batch)
error_test, error_e_test, error_ed_test = test_out
print_str = ""
prop_fmt = " %9.2e %9.2e"
print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0])
print_str += prop_fmt % (np.sqrt(error_ed_test), np.sqrt(error_ed_train))
return print_str
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/loss/loss.py deleted 100644 → 0
View file @ 6b33aeb8
from abc import ABCMeta, abstractmethod
from typing import Tuple, Dict
from deepmd.env import tf
class Loss(metaclass=ABCMeta):
"""The abstract class for the loss function."""
@abstractmethod
def build(self,
learning_rate: tf.Tensor,
natoms: tf.Tensor,
model_dict: Dict[str, tf.Tensor],
label_dict: Dict[str, tf.Tensor],
suffix: str) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:
"""Build the loss function graph.
Parameters
----------
learning_rate : tf.Tensor
learning rate
natoms : tf.Tensor
number of atoms
model_dict : dict[str, tf.Tensor]
A dictionary that maps model keys to tensors
label_dict : dict[str, tf.Tensor]
A dictionary that maps label keys to tensors
suffix : str
suffix
Returns
-------
tf.Tensor
the total squared loss
dict[str, tf.Tensor]
A dictionary that maps loss keys to more loss tensors
"""
@abstractmethod
def eval(self,
sess: tf.Session,
feed_dict: Dict[tf.placeholder, tf.Tensor],
natoms: tf.Tensor) -> dict:
"""Eval the loss function.
Parameters
----------
sess : tf.Session
TensorFlow session
feed_dict : dict[tf.placeholder, tf.Tensor]
A dictionary that maps graph elements to values
natoms : tf.Tensor
number of atoms
Returns
-------
dict
A dictionary that maps keys to values. It
should contain key `natoms`
"""
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/loss/tensor.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from deepmd.env import tf
from deepmd.common import ClassArg, add_data_requirement
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import global_cvt_2_ener_float
from deepmd.utils.sess import run_sess
from .loss import Loss
class TensorLoss(Loss) :
"""
Loss function for tensorial properties.
"""
def __init__ (self, jdata, **kwarg) :
model = kwarg.get('model', None)
if model is not None:
self.type_sel = model.get_sel_type()
else:
self.type_sel = None
self.tensor_name = kwarg['tensor_name']
self.tensor_size = kwarg['tensor_size']
self.label_name = kwarg['label_name']
if jdata is not None:
self.scale = jdata.get('scale', 1.0)
else:
self.scale = 1.0
# YHT: added for global / local dipole combination
assert jdata is not None, "Please provide loss parameters!"
# YWolfeee: modify, use pref / pref_atomic, instead of pref_weight / pref_atomic_weight
self.local_weight = jdata.get('pref_atomic', None)
self.global_weight = jdata.get('pref', None)
assert (self.local_weight is not None and self.global_weight is not None), "Both `pref` and `pref_atomic` should be provided."
assert self.local_weight >= 0.0 and self.global_weight >= 0.0, "Can not assign negative weight to `pref` and `pref_atomic`"
assert (self.local_weight >0.0) or (self.global_weight>0.0), AssertionError('Can not assian zero weight both to `pref` and `pref_atomic`')
# data required
add_data_requirement("atomic_" + self.label_name,
self.tensor_size,
atomic=True,
must=False,
high_prec=False,
type_sel = self.type_sel)
add_data_requirement(self.label_name,
self.tensor_size,
atomic=False,
must=False,
high_prec=False,
type_sel = self.type_sel)
def build (self,
learning_rate,
natoms,
model_dict,
label_dict,
suffix):
polar_hat = label_dict[self.label_name]
atomic_polar_hat = label_dict["atomic_" + self.label_name]
polar = tf.reshape(model_dict[self.tensor_name], [-1])
find_global = label_dict['find_' + self.label_name]
find_atomic = label_dict['find_atomic_' + self.label_name]
# YHT: added for global / local dipole combination
l2_loss = global_cvt_2_tf_float(0.0)
more_loss = {
"local_loss":global_cvt_2_tf_float(0.0),
"global_loss":global_cvt_2_tf_float(0.0)
}
if self.local_weight > 0.0:
local_loss = global_cvt_2_tf_float(find_atomic) * tf.reduce_mean( tf.square(self.scale*(polar - atomic_polar_hat)), name='l2_'+suffix)
more_loss['local_loss'] = local_loss
l2_loss += self.local_weight * local_loss
self.l2_loss_local_summary = tf.summary.scalar('l2_local_loss',
tf.sqrt(more_loss['local_loss']))
if self.global_weight > 0.0: # Need global loss
atoms = 0
if self.type_sel is not None:
for w in self.type_sel:
atoms += natoms[2+w]
else:
atoms = natoms[0]
nframes = tf.shape(polar)[0] // self.tensor_size // atoms
# get global results
global_polar = tf.reshape(tf.reduce_sum(tf.reshape(
polar, [nframes, -1, self.tensor_size]), axis=1),[-1])
#if self.atomic: # If label is local, however
# global_polar_hat = tf.reshape(tf.reduce_sum(tf.reshape(
# polar_hat, [nframes, -1, self.tensor_size]), axis=1),[-1])
#else:
# global_polar_hat = polar_hat
global_loss = global_cvt_2_tf_float(find_global) * tf.reduce_mean( tf.square(self.scale*(global_polar - polar_hat)), name='l2_'+suffix)
more_loss['global_loss'] = global_loss
self.l2_loss_global_summary = tf.summary.scalar('l2_global_loss',
tf.sqrt(more_loss['global_loss']) / global_cvt_2_tf_float(atoms))
# YWolfeee: should only consider atoms with dipole, i.e. atoms
# atom_norm = 1./ global_cvt_2_tf_float(natoms[0])
atom_norm = 1./ global_cvt_2_tf_float(atoms)
global_loss *= atom_norm
l2_loss += self.global_weight * global_loss
self.l2_more = more_loss
self.l2_l = l2_loss
self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
return l2_loss, more_loss
def eval(self, sess, feed_dict, natoms):
atoms = 0
if self.type_sel is not None:
for w in self.type_sel:
atoms += natoms[2+w]
else:
atoms = natoms[0]
run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']]
error, error_lc, error_gl = run_sess(sess, run_data, feed_dict=feed_dict)
results = {"natoms": atoms, "rmse": np.sqrt(error)}
if self.local_weight > 0.0:
results["rmse_lc"] = np.sqrt(error_lc)
if self.global_weight > 0.0:
results["rmse_gl"] = np.sqrt(error_gl) / atoms
return results
def print_header(self): # depreciated
prop_fmt = ' %11s %11s'
print_str = ''
print_str += prop_fmt % ('rmse_tst', 'rmse_trn')
if self.local_weight > 0.0:
print_str += prop_fmt % ('rmse_lc_tst', 'rmse_lc_trn')
if self.global_weight > 0.0:
print_str += prop_fmt % ('rmse_gl_tst', 'rmse_gl_trn')
return print_str
def print_on_training(self,
tb_writer,
cur_batch,
sess,
natoms,
feed_dict_test,
feed_dict_batch) : # depreciated
# YHT: added to calculate the atoms number
atoms = 0
if self.type_sel is not None:
for w in self.type_sel:
atoms += natoms[2+w]
else:
atoms = natoms[0]
run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']]
summary_list = [self.l2_loss_summary]
if self.local_weight > 0.0:
summary_list.append(self.l2_loss_local_summary)
if self.global_weight > 0.0:
summary_list.append(self.l2_loss_global_summary)
# first train data
error_train = run_sess(sess, run_data, feed_dict=feed_dict_batch)
# than test data, if tensorboard log writter is present, commpute summary
# and write tensorboard logs
if tb_writer:
#summary_merged_op = tf.summary.merge([self.l2_loss_summary])
summary_merged_op = tf.summary.merge(summary_list)
run_data.insert(0, summary_merged_op)
test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
if tb_writer:
summary = test_out.pop(0)
tb_writer.add_summary(summary, cur_batch)
error_test = test_out
print_str = ""
prop_fmt = " %11.2e %11.2e"
print_str += prop_fmt % (np.sqrt(error_test[0]), np.sqrt(error_train[0]))
if self.local_weight > 0.0:
print_str += prop_fmt % (np.sqrt(error_test[1]), np.sqrt(error_train[1]) )
if self.global_weight > 0.0:
print_str += prop_fmt % (np.sqrt(error_test[2])/atoms, np.sqrt(error_train[2])/atoms)
return print_str
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/model/__init__.py deleted 100644 → 0
View file @ 6b33aeb8
from .ener import EnerModel
from .tensor import WFCModel
from .tensor import DipoleModel
from .tensor import PolarModel
from .tensor import GlobalPolarModel
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/model/ener.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.utils.pair_tab import PairTab
from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph
from deepmd.utils.errors import GraphWithoutTensorError
from deepmd.common import ClassArg
from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION, GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import op_module
from .model import Model
from .model_stat import make_stat_input, merge_sys_stat
class EnerModel(Model) :
"""Energy model.
Parameters
----------
descrpt
Descriptor
fitting
Fitting net
type_map
Mapping atom type to the name (str) of the type.
For example `type_map[1]` gives the name of the type 1.
data_stat_nbatch
Number of frames used for data statistic
data_stat_protect
Protect parameter for atomic energy regression
use_srtab
The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.
smin_alpha
The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided.
sw_rmin
The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided.
sw_rmin
The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided.
"""
model_type = 'ener'
def __init__ (
self,
descrpt,
fitting,
typeebd = None,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2,
use_srtab : str = None,
smin_alpha : float = None,
sw_rmin : float = None,
sw_rmax : float = None
) -> None:
"""
Constructor
"""
# descriptor
self.descrpt = descrpt
self.rcut = self.descrpt.get_rcut()
self.ntypes = self.descrpt.get_ntypes()
# fitting
self.fitting = fitting
self.numb_fparam = self.fitting.get_numb_fparam()
# type embedding
self.typeebd = typeebd
# other inputs
if type_map is None:
self.type_map = []
else:
self.type_map = type_map
self.data_stat_nbatch = data_stat_nbatch
self.data_stat_protect = data_stat_protect
self.srtab_name = use_srtab
if self.srtab_name is not None :
self.srtab = PairTab(self.srtab_name)
self.smin_alpha = smin_alpha
self.sw_rmin = sw_rmin
self.sw_rmax = sw_rmax
else :
self.srtab = None
def get_rcut (self) :
return self.rcut
def get_ntypes (self) :
return self.ntypes
def get_type_map (self) :
return self.type_map
def data_stat(self, data):
all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False)
m_all_stat = merge_sys_stat(all_stat)
self._compute_input_stat(m_all_stat, protection=self.data_stat_protect, mixed_type=data.mixed_type)
self._compute_output_stat(all_stat, mixed_type=data.mixed_type)
# self.bias_atom_e = data.compute_energy_shift(self.rcond)
def _compute_input_stat (self, all_stat, protection=1e-2, mixed_type=False):
if mixed_type:
self.descrpt.compute_input_stats(all_stat['coord'],
all_stat['box'],
all_stat['type'],
all_stat['natoms_vec'],
all_stat['default_mesh'],
all_stat,
mixed_type,
all_stat['real_natoms_vec'])
else:
self.descrpt.compute_input_stats(all_stat['coord'],
all_stat['box'],
all_stat['type'],
all_stat['natoms_vec'],
all_stat['default_mesh'],
all_stat)
self.fitting.compute_input_stats(all_stat, protection=protection)
def _compute_output_stat (self, all_stat, mixed_type=False):
if mixed_type:
self.fitting.compute_output_stats(all_stat, mixed_type=mixed_type)
else:
self.fitting.compute_output_stats(all_stat)
def build (self,
coord_,
atype_,
natoms,
box,
mesh,
input_dict,
frz_model = None,
suffix = '',
reuse = None):
if input_dict is None:
input_dict = {}
with tf.variable_scope('model_attr' + suffix, reuse = reuse) :
t_tmap = tf.constant(' '.join(self.type_map),
name = 'tmap',
dtype = tf.string)
t_mt = tf.constant(self.model_type,
name = 'model_type',
dtype = tf.string)
t_ver = tf.constant(MODEL_VERSION,
name = 'model_version',
dtype = tf.string)
if self.srtab is not None :
tab_info, tab_data = self.srtab.get()
self.tab_info = tf.get_variable('t_tab_info',
tab_info.shape,
dtype = tf.float64,
trainable = False,
initializer = tf.constant_initializer(tab_info, dtype = tf.float64))
self.tab_data = tf.get_variable('t_tab_data',
tab_data.shape,
dtype = tf.float64,
trainable = False,
initializer = tf.constant_initializer(tab_data, dtype = tf.float64))
coord = tf.reshape (coord_, [-1, natoms[1] * 3])
atype = tf.reshape (atype_, [-1, natoms[1]])
input_dict['nframes'] = tf.shape(coord)[0]
# type embedding if any
if self.typeebd is not None:
type_embedding = self.typeebd.build(
self.ntypes,
reuse = reuse,
suffix = suffix,
)
input_dict['type_embedding'] = type_embedding
input_dict['atype'] = atype_
if frz_model == None:
dout \
= self.descrpt.build(coord_,
atype_,
natoms,
box,
mesh,
input_dict,
suffix = suffix,
reuse = reuse)
dout = tf.identity(dout, name='o_descriptor')
else:
tf.constant(self.rcut,
name = 'descrpt_attr/rcut',
dtype = GLOBAL_TF_FLOAT_PRECISION)
tf.constant(self.ntypes,
name = 'descrpt_attr/ntypes',
dtype = tf.int32)
feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh)
return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0']
imported_tensors \
= self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements)
dout = imported_tensors[-1]
self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1])
if self.srtab is not None :
nlist, rij, sel_a, sel_r = self.descrpt.get_nlist()
nnei_a = np.cumsum(sel_a)[-1]
nnei_r = np.cumsum(sel_r)[-1]
atom_ener = self.fitting.build (dout,
natoms,
input_dict,
reuse = reuse,
suffix = suffix)
self.atom_ener = atom_ener
if self.srtab is not None :
sw_lambda, sw_deriv \
= op_module.soft_min_switch(atype,
rij,
nlist,
natoms,
sel_a = sel_a,
sel_r = sel_r,
alpha = self.smin_alpha,
rmin = self.sw_rmin,
rmax = self.sw_rmax)
inv_sw_lambda = 1.0 - sw_lambda
# NOTICE:
# atom energy is not scaled,
# force and virial are scaled
tab_atom_ener, tab_force, tab_atom_virial \
= op_module.pair_tab(self.tab_info,
self.tab_data,
atype,
rij,
nlist,
natoms,
sw_lambda,
sel_a = sel_a,
sel_r = sel_r)
energy_diff = tab_atom_ener - tf.reshape(atom_ener, [-1, natoms[0]])
tab_atom_ener = tf.reshape(sw_lambda, [-1]) * tf.reshape(tab_atom_ener, [-1])
atom_ener = tf.reshape(inv_sw_lambda, [-1]) * atom_ener
energy_raw = tab_atom_ener + atom_ener
else :
energy_raw = atom_ener
energy_raw = tf.reshape(energy_raw, [-1, natoms[0]], name = 'o_atom_energy'+suffix)
energy = tf.reduce_sum(global_cvt_2_ener_float(energy_raw), axis=1, name='o_energy'+suffix)
force, virial, atom_virial \
= self.descrpt.prod_force_virial (atom_ener, natoms)
if self.srtab is not None :
sw_force \
= op_module.soft_min_force(energy_diff,
sw_deriv,
nlist,
natoms,
n_a_sel = nnei_a,
n_r_sel = nnei_r)
force = force + sw_force + tab_force
force = tf.reshape (force, [-1, 3 * natoms[1]], name = "o_force"+suffix)
if self.srtab is not None :
sw_virial, sw_atom_virial \
= op_module.soft_min_virial (energy_diff,
sw_deriv,
rij,
nlist,
natoms,
n_a_sel = nnei_a,
n_r_sel = nnei_r)
atom_virial = atom_virial + sw_atom_virial + tab_atom_virial
virial = virial + sw_virial \
+ tf.reduce_sum(tf.reshape(tab_atom_virial, [-1, natoms[1], 9]), axis = 1)
virial = tf.reshape (virial, [-1, 9], name = "o_virial"+suffix)
atom_virial = tf.reshape (atom_virial, [-1, 9 * natoms[1]], name = "o_atom_virial"+suffix)
model_dict = {}
model_dict['energy'] = energy
model_dict['force'] = force
model_dict['virial'] = virial
model_dict['atom_ener'] = energy_raw
model_dict['atom_virial'] = atom_virial
model_dict['coord'] = coord
model_dict['atype'] = atype
return model_dict
def _import_graph_def_from_frz_model(self, frz_model, feed_dict, return_elements):
graph, graph_def = load_graph_def(frz_model)
return tf.import_graph_def(graph_def, input_map = feed_dict, return_elements = return_elements, name = "")
def init_variables(self,
graph : tf.Graph,
graph_def : tf.GraphDef,
model_type : str = "original_model",
suffix : str = "",
) -> None:
"""
Init the embedding net variables with the given frozen model
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
model_type : str
the type of the model
suffix : str
suffix to name scope
"""
# self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch...
# initialize fitting net with the given compressed frozen model
if model_type == 'original_model':
self.descrpt.init_variables(graph, graph_def, suffix=suffix)
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("original_model", name = 'model_type', dtype = tf.string)
elif model_type == 'compressed_model':
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
else:
raise RuntimeError("Unknown model type %s" % model_type)
if self.typeebd is not None:
self.typeebd.init_variables(graph, graph_def, suffix=suffix)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/model/model.py deleted 100644 → 0
View file @ 6b33aeb8
from deepmd.env import tf
class Model:
def init_variables(self,
graph : tf.Graph,
graph_def : tf.GraphDef,
model_type : str = "original_model",
suffix : str = "",
) -> None:
"""
Init the embedding net variables with the given frozen model
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
model_type : str
the type of the model
suffix : str
suffix to name scope
"""
raise RuntimeError("The 'dp train init-frz-model' command do not support this model!")
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/model/model_stat.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from collections import defaultdict
def _make_all_stat_ref(data, nbatches):
all_stat = defaultdict(list)
for ii in range(data.get_nsystems()) :
for jj in range(nbatches) :
stat_data = data.get_batch (sys_idx = ii)
for dd in stat_data:
if dd == "natoms_vec":
stat_data[dd] = stat_data[dd].astype(np.int32)
all_stat[dd].append(stat_data[dd])
return all_stat
def make_stat_input(data, nbatches, merge_sys = True):
"""
pack data for statistics
Parameters
----------
data:
The data
merge_sys: bool (True)
Merge system data
Returns
-------
all_stat:
A dictionary of list of list storing data for stat.
if merge_sys == False data can be accessed by
all_stat[key][sys_idx][batch_idx][frame_idx]
else merge_sys == True can be accessed by
all_stat[key][batch_idx][frame_idx]
"""
all_stat = defaultdict(list)
for ii in range(data.get_nsystems()) :
sys_stat = defaultdict(list)
for jj in range(nbatches) :
stat_data = data.get_batch (sys_idx = ii)
for dd in stat_data:
if dd == "natoms_vec":
stat_data[dd] = stat_data[dd].astype(np.int32)
sys_stat[dd].append(stat_data[dd])
for dd in sys_stat:
if merge_sys:
for bb in sys_stat[dd]:
all_stat[dd].append(bb)
else:
all_stat[dd].append(sys_stat[dd])
return all_stat
def merge_sys_stat(all_stat):
first_key = list(all_stat.keys())[0]
nsys = len(all_stat[first_key])
ret = defaultdict(list)
for ii in range(nsys):
for dd in all_stat:
for bb in all_stat[dd][ii]:
ret[dd].append(bb)
return ret
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/model/tensor.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from typing import Tuple, List
from deepmd.env import tf
from deepmd.common import ClassArg
from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION, GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import op_module
from deepmd.utils.graph import load_graph_def
from .model import Model
from .model_stat import make_stat_input, merge_sys_stat
class TensorModel(Model) :
"""Tensor model.
Parameters
----------
tensor_name
Name of the tensor.
descrpt
Descriptor
fitting
Fitting net
type_map
Mapping atom type to the name (str) of the type.
For example `type_map[1]` gives the name of the type 1.
data_stat_nbatch
Number of frames used for data statistic
data_stat_protect
Protect parameter for atomic energy regression
"""
def __init__ (
self,
tensor_name : str,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2,
)->None:
"""
Constructor
"""
self.model_type = tensor_name
# descriptor
self.descrpt = descrpt
self.rcut = self.descrpt.get_rcut()
self.ntypes = self.descrpt.get_ntypes()
# fitting
self.fitting = fitting
# other params
if type_map is None:
self.type_map = []
else:
self.type_map = type_map
self.data_stat_nbatch = data_stat_nbatch
self.data_stat_protect = data_stat_protect
def get_rcut (self) :
return self.rcut
def get_ntypes (self) :
return self.ntypes
def get_type_map (self) :
return self.type_map
def get_sel_type(self):
return self.fitting.get_sel_type()
def get_out_size (self) :
return self.fitting.get_out_size()
def data_stat(self, data):
all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False)
m_all_stat = merge_sys_stat(all_stat)
self._compute_input_stat (m_all_stat, protection = self.data_stat_protect)
self._compute_output_stat(all_stat)
def _compute_input_stat(self, all_stat, protection = 1e-2) :
self.descrpt.compute_input_stats(all_stat['coord'],
all_stat['box'],
all_stat['type'],
all_stat['natoms_vec'],
all_stat['default_mesh'],
all_stat)
if hasattr(self.fitting, 'compute_input_stats'):
self.fitting.compute_input_stats(all_stat, protection = protection)
def _compute_output_stat (self, all_stat) :
if hasattr(self.fitting, 'compute_output_stats'):
self.fitting.compute_output_stats(all_stat)
def build (self,
coord_,
atype_,
natoms,
box,
mesh,
input_dict,
frz_model = None,
suffix = '',
reuse = None):
with tf.variable_scope('model_attr' + suffix, reuse = reuse) :
t_tmap = tf.constant(' '.join(self.type_map),
name = 'tmap',
dtype = tf.string)
t_st = tf.constant(self.get_sel_type(),
name = 'sel_type',
dtype = tf.int32)
t_mt = tf.constant(self.model_type,
name = 'model_type',
dtype = tf.string)
t_ver = tf.constant(MODEL_VERSION,
name = 'model_version',
dtype = tf.string)
t_od = tf.constant(self.get_out_size(),
name = 'output_dim',
dtype = tf.int32)
natomsel = sum(natoms[2+type_i] for type_i in self.get_sel_type())
nout = self.get_out_size()
if frz_model == None:
dout \
= self.descrpt.build(coord_,
atype_,
natoms,
box,
mesh,
input_dict,
suffix = suffix,
reuse = reuse)
dout = tf.identity(dout, name='o_descriptor')
else:
tf.constant(self.rcut,
name = 'descrpt_attr/rcut',
dtype = GLOBAL_TF_FLOAT_PRECISION)
tf.constant(self.ntypes,
name = 'descrpt_attr/ntypes',
dtype = tf.int32)
feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh)
return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0']
imported_tensors \
= self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements)
dout = imported_tensors[-1]
self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1])
rot_mat = self.descrpt.get_rot_mat()
rot_mat = tf.identity(rot_mat, name = 'o_rot_mat'+suffix)
output = self.fitting.build (dout,
rot_mat,
natoms,
reuse = reuse,
suffix = suffix)
framesize = nout if "global" in self.model_type else natomsel * nout
output = tf.reshape(output, [-1, framesize], name = 'o_' + self.model_type + suffix)
model_dict = {self.model_type: output}
if "global" not in self.model_type:
gname = "global_"+self.model_type
atom_out = tf.reshape(output, [-1, natomsel, nout])
global_out = tf.reduce_sum(atom_out, axis=1)
global_out = tf.reshape(global_out, [-1, nout], name="o_" + gname + suffix)
out_cpnts = tf.split(atom_out, nout, axis=-1)
force_cpnts = []
virial_cpnts = []
atom_virial_cpnts = []
for out_i in out_cpnts:
force_i, virial_i, atom_virial_i \
= self.descrpt.prod_force_virial(out_i, natoms)
force_cpnts.append (tf.reshape(force_i, [-1, 3*natoms[1]]))
virial_cpnts.append (tf.reshape(virial_i, [-1, 9]))
atom_virial_cpnts.append(tf.reshape(atom_virial_i, [-1, 9*natoms[1]]))
# [nframe x nout x (natom x 3)]
force = tf.concat(force_cpnts, axis=1, name="o_force" + suffix)
# [nframe x nout x 9]
virial = tf.concat(virial_cpnts, axis=1, name="o_virial" + suffix)
# [nframe x nout x (natom x 9)]
atom_virial = tf.concat(atom_virial_cpnts, axis=1, name="o_atom_virial" + suffix)
model_dict[gname] = global_out
model_dict["force"] = force
model_dict["virial"] = virial
model_dict["atom_virial"] = atom_virial
return model_dict
def _import_graph_def_from_frz_model(self, frz_model, feed_dict, return_elements):
graph, graph_def = load_graph_def(frz_model)
return tf.import_graph_def(graph_def, input_map = feed_dict, return_elements = return_elements, name = "")
def init_variables(self,
graph : tf.Graph,
graph_def : tf.GraphDef,
model_type : str = "original_model",
suffix : str = "",
) -> None:
"""
Init the embedding net variables with the given frozen model
Parameters
----------
graph : tf.Graph
The input frozen model graph
graph_def : tf.GraphDef
The input frozen model graph_def
model_type : str
the type of the model
suffix : str
suffix to name scope
"""
if model_type == 'original_model':
self.descrpt.init_variables(graph, graph_def, suffix=suffix)
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("original_model", name = 'model_type', dtype = tf.string)
elif model_type == 'compressed_model':
self.fitting.init_variables(graph, graph_def, suffix=suffix)
tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
else:
raise RuntimeError("Unknown model type %s" % model_type)
class WFCModel(TensorModel):
def __init__(
self,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2
) -> None:
TensorModel.__init__(self, 'wfc', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
class DipoleModel(TensorModel):
def __init__(
self,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2
) -> None:
TensorModel.__init__(self, 'dipole', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
class PolarModel(TensorModel):
def __init__(
self,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2
) -> None:
TensorModel.__init__(self, 'polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
class GlobalPolarModel(TensorModel):
def __init__(
self,
descrpt,
fitting,
type_map : List[str] = None,
data_stat_nbatch : int = 10,
data_stat_protect : float = 1e-2
) -> None:
TensorModel.__init__(self, 'global_polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/nvnmd/__init__.py deleted 100644 → 0
View file @ 6b33aeb8
from . import data, descriptor, entrypoints, fit, utils
__all__ = [
"data",
"descriptor",
"entrypoints",
"fit",
"utils",
]
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/nvnmd/data/__init__.py deleted 100644 → 0
View file @ 6b33aeb8
"""
nvnmd.data
==========
Provides
1. hardware configuration
2. default input script
3. title and citation
Data
----
jdata_sys
action configuration
jdata_config
hardware configuration
dscp
descriptor configuration
fitn
fitting network configuration
size
ram capacity
ctrl
control flag, such as Time Division Multiplexing (TDM)
nbit
number of bits of fixed-point number
jdata_config_16 (disable)
difference with configure fitting size as 16
jdata_config_32 (disable)
difference with configure fitting size as 32
jdata_config_64 (disable)
difference with configure fitting size as 64
jdata_config_128 (default)
difference with configure fitting size as 128
jdata_configs
all configure of jdata_config{nfit_node}
jdata_deepmd_input
default input script for nvnmd training
NVNMD_WELCOME
nvnmd title when logging
NVNMD_CITATION
citation of nvnmd
"""
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/nvnmd/data/data.py deleted 100644 → 0
View file @ 6b33aeb8
jdata_sys = {
"debug": False
}
jdata_config = {
"dscp": {
"sel": [60, 60],
"rcut": 6.0,
"rcut_smth": 0.5,
"neuron": [8, 16, 32],
"resnet_dt": False,
"axis_neuron": 4,
"type_one_side": True,
"NI": 128,
"rc_lim": 0.5,
"M1": "neuron[-1]",
"M2": "axis_neuron",
"SEL": [60, 60, 0, 0],
"NNODE_FEAS": "(1, neuron)",
"nlayer_fea": "len(neuron)",
"same_net": "type_one_side",
"NIDP": "sum(sel)",
"NIX": "2^ceil(ln2(NIDP/1.5))",
"ntype": "len(sel)",
"ntypex": "same_net ? 1: ntype",
"ntypex_max": 1,
"ntype_max": 4
},
"fitn": {
"neuron": [32, 32, 32],
"resnet_dt": False,
"NNODE_FITS": "(M1*M2, neuron, 1)",
"nlayer_fit": "len(neuron)+1",
"NLAYER": "nlayer_fit"
},
"size": {
"NTYPE_MAX": 4,
"NSPU": 4096,
"MSPU": 32768,
"Na": "NSPU",
"NaX": "MSPU"
},
"ctrl": {
"NSTDM": 16,
"NSTDM_M1": 16,
"NSTDM_M2": 1,
"NSADV": "NSTDM+1",
"NSEL": "NSTDM*ntype_max",
"NSTDM_M1X": 4,
"NSTEP_DELAY": 20,
"MAX_FANOUT": 30
},
"nbit": {
"NBIT_DATA": 21,
"NBIT_DATA_FL": 13,
"NBIT_LONG_DATA": 32,
"NBIT_LONG_DATA_FL": 24,
"NBIT_DIFF_DATA": 24,
"NBIT_SPE": 2,
"NBIT_CRD": "NBIT_DATA*3",
"NBIT_LST": "ln2(NaX)",
"NBIT_SPE_MAX": 8,
"NBIT_LST_MAX": 16,
"NBIT_ATOM": "NBIT_SPE+NBIT_CRD",
"NBIT_LONG_ATOM": "NBIT_SPE+NBIT_LONG_DATA*3",
"NBIT_RIJ": "NBIT_DATA_FL+5",
"NBIT_FEA_X": 10,
"NBIT_FEA_X_FL": 4,
"NBIT_FEA_X2_FL": 6,
"NBIT_FEA": 18,
"NBIT_FEA_FL": 10,
"NBIT_SHIFT": 4,
"NBIT_DATA2": "NBIT_DATA+NBIT_DATA_FL",
"NBIT_DATA2_FL": "2*NBIT_DATA_FL",
"NBIT_DATA_FEA": "NBIT_DATA+NBIT_FEA_FL",
"NBIT_DATA_FEA_FL": "NBIT_DATA_FL+NBIT_FEA_FL",
"NBIT_FORCE": 32,
"NBIT_FORCE_FL": "2*NBIT_DATA_FL-1",
"NBIT_SUM": "NBIT_DATA_FL+8",
"NBIT_WEIGHT": 18,
"NBIT_WEIGHT_FL": 13,
"NBIT_RAM": 72,
"NBIT_ADDR": 32,
"NBTI_MODEL_HEAD": 32,
"NBIT_TH_LONG_ADD": 30,
"NBIT_ADD": 15,
"RANGE_B": [-100, 100],
"RANGE_W": [-20, 20],
"NCFG": 35,
"NNET": 4920,
"NFEA": 8192
},
"end": ""
}
jdata_config_16 = {
"dscp": {
"neuron": [8, 16, 32],
"axis_neuron": 4,
"NI": 128
},
"fitn": {
"neuron": [16, 16, 16]
},
"ctrl": {
"NSTDM": 16,
"NSTDM_M1": 16,
"NSTDM_M2": 1,
"NSTDM_M1X": 4
}
}
jdata_config_32 = {
"dscp": {
"neuron": [8, 16, 32],
"axis_neuron": 4,
"NI": 128
},
"fitn": {
"neuron": [32, 32, 32]
},
"ctrl": {
"NSTDM": 16,
"NSTDM_M1": 16,
"NSTDM_M2": 1,
"NSTDM_M1X": 4
}
}
jdata_config_64 = {
"dscp": {
"neuron": [8, 16, 32],
"axis_neuron": 4,
"NI": 128
},
"fitn": {
"neuron": [64, 64, 64]
},
"ctrl": {
"NSTDM": 32,
"NSTDM_M1": 32,
"NSTDM_M2": 1,
"NSTDM_M1X": 4
}
}
jdata_config_128 = {
"dscp": {
"neuron": [8, 16, 32],
"axis_neuron": 4,
"NI": 128
},
"fitn": {
"neuron": [128, 128, 128]
},
"ctrl": {
"NSTDM": 32,
"NSTDM_M1": 32,
"NSTDM_M2": 1,
"NSTDM_M1X": 4
}
}
jdata_configs = {
"_16": jdata_config_16,
"_32": jdata_config_32,
"_64": jdata_config_64,
"128": jdata_config_128
}
jdata_deepmd_input = {
"model": {
"descriptor": {
"seed": 1,
"type": "se_a",
"sel": [
60,
60
],
"rcut": 7.0,
"rcut_smth": 0.5,
"neuron": [
8,
16,
32
],
"type_one_side": False,
"axis_neuron": 4,
"resnet_dt": False
},
"fitting_net": {
"seed": 1,
"neuron": [
128,
128,
128
],
"resnet_dt": False
}
},
"nvnmd": {
"net_size": 128,
"config_file": "none",
"weight_file": "none",
"map_file": "none",
"enable": False,
"restore_descriptor": False,
"restore_fitting_net": False,
"quantize_descriptor": False,
"quantize_fitting_net": False
},
"learning_rate": {
"type": "exp",
"decay_steps": 5000,
"start_lr": 0.005,
"stop_lr": 8.257687192506788e-05
},
"loss": {
"start_pref_e": 0.02,
"limit_pref_e": 1,
"start_pref_f": 1000,
"limit_pref_f": 1,
"start_pref_v": 0,
"limit_pref_v": 0
},
"training": {
"seed": 1,
"stop_batch": 10000,
"disp_file": "lcurve.out",
"disp_freq": 100,
"numb_test": 10,
"save_freq": 1000,
"save_ckpt": "model.ckpt",
"disp_training": True,
"time_training": True,
"profiling": False,
"training_data": {
"systems": "dataset",
"set_prefix": "set",
"batch_size": 1
}
}
}
NVNMD_WELCOME = (
" _ _ __ __ _ _ __ __ ____ ",
"| \ | | \ \ / / | \ | | | \/ | | _ \ ",
"| \| | \ \ / / | \| | | |\/| | | | | |",
"| |\ | \ V / | |\ | | | | | | |_| |",
"|_| \_| \_/ |_| \_| |_| |_| |____/ ",
)
NVNMD_CITATION = (
"Please read and cite:",
"Mo et al., npj Comput Mater 8, 107 (2022)",
)
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/nvnmd/descriptor/__init__.py deleted 100644 → 0
View file @ 6b33aeb8
"""
nvnmd.se_a
==========
Provides
1. building descriptor with continuous embedding network
2. building descriptor with quantized embedding network
"""
_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/nvnmd/descriptor/se_a.py deleted 100644 → 0
View file @ 6b33aeb8
import numpy as np
from deepmd.env import tf
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
from deepmd.env import op_module
from deepmd.utils.network import embedding_net
#
from deepmd.nvnmd.utils.config import nvnmd_cfg
from deepmd.nvnmd.utils.network import matmul3_qq
from deepmd.nvnmd.utils.weight import get_normalize, get_rng_s
def build_davg_dstd():
r"""Get the davg and dstd from the dictionary nvnmd_cfg.
The davg and dstd have been obtained by training CNN
"""
davg, dstd = get_normalize(nvnmd_cfg.weight)
return davg, dstd
def build_op_descriptor():
r"""Replace se_a.py/DescrptSeA/build
"""
if nvnmd_cfg.quantize_descriptor:
return op_module.prod_env_mat_a_nvnmd_quantize
else:
return op_module.prod_env_mat_a
def descrpt2r4(inputs, natoms):
r"""Replace :math:`r_{ji} \rightarrow r'_{ji}`
where :math:`r_{ji} = (x_{ji}, y_{ji}, z_{ji})` and
:math:`r'_{ji} = (s_{ji}, \frac{s_{ji} x_{ji}}{r_{ji}}, \frac{s_{ji} y_{ji}}{r_{ji}}, \frac{s_{ji} z_{ji}}{r_{ji}})`
"""
NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
NBIT_FEA_X_FL = nvnmd_cfg.nbit['NBIT_FEA_X_FL']
NBIT_FEA_FL = nvnmd_cfg.nbit['NBIT_FEA_FL']
prec = 1.0 / (2 ** NBIT_FEA_X_FL)
ntypes = nvnmd_cfg.dscp['ntype']
NIDP = nvnmd_cfg.dscp['NIDP']
ndescrpt = NIDP * 4
start_index = 0
# (nf, na*nd)
shape = inputs.get_shape().as_list()
# (nf*na*ni, 4)
inputs_reshape = tf.reshape(inputs, [-1, 4])
with tf.variable_scope('filter_type_all_x', reuse=True):
# u (i.e., r^2)
u = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
with tf.variable_scope('u', reuse=True):
u = op_module.quantize_nvnmd(u, 0, -1, NBIT_DATA_FL, -1)
# print('u:', u)
u = tf.reshape(u, [-1, natoms[0] * NIDP])
# rij
rij = tf.reshape(tf.slice(inputs_reshape, [0, 1], [-1, 3]), [-1, 3])
with tf.variable_scope('rij', reuse=True):
rij = op_module.quantize_nvnmd(rij, 0, NBIT_DATA_FL, -1, -1)
# print('rij:', rij)
s = []
sr = []
for type_i in range(ntypes):
type_input = 0
postfix = f"_t{type_input}_t{type_i}"
u_i = tf.slice(
u,
[0, start_index * NIDP],
[-1, natoms[2 + type_i] * NIDP])
u_i = tf.reshape(u_i, [-1, 1])
#
keys = 's,sr'.split(',')
map_tables = [nvnmd_cfg.map[key + postfix] for key in keys]
map_tables2 = [nvnmd_cfg.map[f"d{key}_dr2" + postfix] for key in keys]
map_outs = []
for ii in range(len(keys)):
map_outs.append(op_module.map_nvnmd(
u_i,
map_tables[ii][0],
map_tables[ii][1] / prec,
map_tables2[ii][0],
map_tables2[ii][1] / prec,
prec, NBIT_FEA_FL))
s_i, sr_i = map_outs
s_i = tf.reshape(s_i, [-1, natoms[2 + type_i] * NIDP])
sr_i = tf.reshape(sr_i, [-1, natoms[2 + type_i] * NIDP])
s.append(s_i)
sr.append(sr_i)
start_index += natoms[2 + type_i]
s = tf.concat(s, axis=1)
sr = tf.concat(sr, axis=1)
with tf.variable_scope('s', reuse=True):
s = op_module.quantize_nvnmd(s, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
with tf.variable_scope('sr', reuse=True):
sr = op_module.quantize_nvnmd(sr, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
s = tf.reshape(s, [-1, 1])
sr = tf.reshape(sr, [-1, 1])
# R2R4
Rs = s
Rxyz = sr * rij
with tf.variable_scope('Rxyz', reuse=True):
Rxyz = op_module.quantize_nvnmd(Rxyz, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
R4 = tf.concat([Rs, Rxyz], axis=1)
R4 = tf.reshape(R4, [-1, NIDP, 4])
inputs_reshape = R4
inputs_reshape = tf.reshape(inputs_reshape, [-1, ndescrpt])
return inputs_reshape
def filter_lower_R42GR(
type_i,
type_input,
inputs_i,
is_exclude,
activation_fn,
bavg,
stddev,
trainable,
suffix,
seed,
seed_shift,
uniform_seed,
filter_neuron,
filter_precision,
filter_resnet_dt,
embedding_net_variables):
r"""Replace se_a.py/DescrptSeA/_filter_lower
"""
shape_i = inputs_i.get_shape().as_list()
inputs_reshape = tf.reshape(inputs_i, [-1, 4])
natom = tf.shape(inputs_i)[0]
M1 = nvnmd_cfg.dscp['M1']
NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
NBIT_FEA_X_FL = nvnmd_cfg.nbit['NBIT_FEA_X_FL']
NBIT_FEA_X2_FL = nvnmd_cfg.nbit['NBIT_FEA_X2_FL']
NBIT_FEA_FL = nvnmd_cfg.nbit['NBIT_FEA_FL']
prec = 1.0 / (2 ** NBIT_FEA_X2_FL)
type_input = 0 if (type_input < 0) else type_input
postfix = f"_t{type_input}_t{type_i}"
if (nvnmd_cfg.quantize_descriptor):
s_min, smax = get_rng_s(nvnmd_cfg.weight)
s_min = -2.0
# s_min = np.floor(s_min)
s = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
s = op_module.quantize_nvnmd(s, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
# G
keys = 'G'.split(',')
map_tables = [nvnmd_cfg.map[key + postfix] for key in keys]
map_tables2 = [nvnmd_cfg.map[f"d{key}_ds" + postfix] for key in keys]
map_outs = []
for ii in range(len(keys)):
with tf.variable_scope(keys[ii], reuse=True):
map_outs.append(op_module.map_nvnmd(
s - s_min,
map_tables[ii][0], map_tables[ii][1] / prec,
map_tables2[ii][0], map_tables2[ii][1] / prec,
prec, NBIT_FEA_FL))
map_outs[ii] = op_module.quantize_nvnmd(map_outs[ii], 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
G = map_outs
# G
xyz_scatter = G
xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, M1))
# GR
inputs_reshape = tf.reshape(inputs_reshape, [-1, shape_i[1] // 4, 4])
GR = matmul3_qq(tf.transpose(inputs_reshape, [0, 2, 1]), xyz_scatter, -1)
GR = tf.reshape(GR, [-1, 4 * M1])
return GR
else:
xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
if nvnmd_cfg.restore_descriptor:
trainable = False
embedding_net_variables = {}
for key in nvnmd_cfg.weight.keys():
if 'filter_type' in key:
key2 = key.replace('.', '/')
embedding_net_variables[key2] = nvnmd_cfg.weight[key]
if (not is_exclude):
xyz_scatter = embedding_net(
xyz_scatter,
filter_neuron,
filter_precision,
activation_fn=activation_fn,
resnet_dt=filter_resnet_dt,
name_suffix=suffix,
stddev=stddev,
bavg=bavg,
seed=seed,
trainable=trainable,
uniform_seed=uniform_seed,
initial_variables=embedding_net_variables)
if (not uniform_seed) and (seed is not None):
seed += seed_shift
else:
# we can safely return the final xyz_scatter filled with zero directly
return tf.cast(tf.fill((natom, 4, M1), 0.), GLOBAL_TF_FLOAT_PRECISION)
# natom x nei_type_i x out_size
xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, M1))
# When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below
# [588 24] -> [588 6 4] correct
# but if sel is zero
# [588 0] -> [147 0 4] incorrect; the correct one is [588 0 4]
# So we need to explicitly assign the shape to tf.shape(inputs_i)[0] instead of -1
return tf.matmul(tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), xyz_scatter, transpose_a=True)
def filter_GR2D(xyz_scatter_1):
r"""Replace se_a.py/_filter
"""
NIX = nvnmd_cfg.dscp['NIX']
NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
M1 = nvnmd_cfg.dscp['M1']
M2 = nvnmd_cfg.dscp['M2']
if (nvnmd_cfg.quantize_descriptor):
xyz_scatter_1 = tf.reshape(xyz_scatter_1, [-1, 4 * M1])
# fix the number of bits of gradient
xyz_scatter_1 = op_module.quantize_nvnmd(xyz_scatter_1, 0, -1, NBIT_DATA_FL, -1)
xyz_scatter_1 = xyz_scatter_1 * (1.0 / NIX)
with tf.variable_scope('GR', reuse=True):
xyz_scatter_1 = op_module.quantize_nvnmd(xyz_scatter_1, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
xyz_scatter_1 = tf.reshape(xyz_scatter_1, [-1, 4, M1])
# natom x 4 x outputs_size_2
xyz_scatter_2 = xyz_scatter_1
# natom x 3 x outputs_size_1
qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1])
# natom x outputs_size_2 x 3
qmat = tf.transpose(qmat, perm=[0, 2, 1])
# D': natom x outputs_size x outputs_size_2
result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True)
# D': natom x (outputs_size x outputs_size_2)
result = tf.reshape(result, [-1, M1 * M1])
#
index_subset = []
for ii in range(M1):
for jj in range(ii, ii + M2):
index_subset.append((ii * M1) + (jj % M1))
index_subset = tf.constant(np.int32(np.array(index_subset)))
result = tf.gather(result, index_subset, axis=1)
with tf.variable_scope('d', reuse=True):
result = op_module.quantize_nvnmd(result, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
else:
# natom x 4 x outputs_size
xyz_scatter_1 = xyz_scatter_1 * (1.0 / NIX)
# natom x 4 x outputs_size_2
# xyz_scatter_2 = tf.slice(xyz_scatter_1, [0,0,0],[-1,-1,outputs_size_2])
xyz_scatter_2 = xyz_scatter_1
# natom x 3 x outputs_size_1
qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1])
# natom x outputs_size_1 x 3
qmat = tf.transpose(qmat, perm=[0, 2, 1])
# natom x outputs_size x outputs_size_2
result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True)
# natom x (outputs_size x outputs_size_2)
# result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]])
result = tf.reshape(result, [-1, M1 * M1])
#
index_subset = []
for ii in range(M1):
for jj in range(ii, ii + M2):
index_subset.append((ii * M1) + (jj % M1))
index_subset = tf.constant(np.int32(np.array(index_subset)))
result = tf.gather(result, index_subset, axis=1)
return result, qmat
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment