From 2fc1eaec9a8d7212fff0c8b14a96dcb2c814a595 Mon Sep 17 00:00:00 2001 From: jxxiaoshaoye <1700017801@pku.edu.cn> Date: Tue, 27 Apr 2021 21:00:53 +0800 Subject: [PATCH 1/2] add type embedding --- deepmd/descriptor/__init__.py | 1 + deepmd/descriptor/se_a_type.py | 807 +++++++++++++++++++++++++ deepmd/fit/ener.py | 128 +++- deepmd/model/ener.py | 24 +- deepmd/train/trainer.py | 10 +- deepmd/utils/argcheck.py | 39 +- deepmd/utils/network.py | 291 +++++++++ deepmd/utils/type_embed.py | 66 ++ examples/water/se_a_type/input.json | 69 +++ source/tests/e.out | 1 + source/tests/f.out | 1 + source/tests/system/set.000/aparam.npy | Bin 0 -> 224 bytes source/tests/system/set.000/box.npy | Bin 0 -> 200 bytes source/tests/system/set.000/coord.npy | Bin 0 -> 272 bytes source/tests/system/set.000/energy.npy | Bin 0 -> 136 bytes source/tests/system/set.000/force.npy | Bin 0 -> 272 bytes source/tests/system/set.000/fparam.npy | Bin 0 -> 144 bytes source/tests/system/type.raw | 6 + source/tests/system/type_map.raw | 2 + source/tests/v.out | 1 + 20 files changed, 1436 insertions(+), 10 deletions(-) create mode 100644 deepmd/descriptor/se_a_type.py create mode 100644 deepmd/utils/type_embed.py create mode 100644 examples/water/se_a_type/input.json create mode 100644 source/tests/e.out create mode 100644 source/tests/f.out create mode 100644 source/tests/system/set.000/aparam.npy create mode 100644 source/tests/system/set.000/box.npy create mode 100644 source/tests/system/set.000/coord.npy create mode 100644 source/tests/system/set.000/energy.npy create mode 100644 source/tests/system/set.000/force.npy create mode 100644 source/tests/system/set.000/fparam.npy create mode 100644 source/tests/system/type.raw create mode 100644 source/tests/system/type_map.raw create mode 100644 source/tests/v.out diff --git a/deepmd/descriptor/__init__.py b/deepmd/descriptor/__init__.py index 7c3d910091..e736ef55ae 100644 --- a/deepmd/descriptor/__init__.py +++ b/deepmd/descriptor/__init__.py @@ -4,6 +4,7 @@ from .se_ar import DescrptSeAR from .se_t import DescrptSeT from .se_a_ebd import DescrptSeAEbd +from .se_a_type import DescrptSeAEbd_type from .se_a_ef import DescrptSeAEf from .se_a_ef import DescrptSeAEfLower from .loc_frame import DescrptLocFrame diff --git a/deepmd/descriptor/se_a_type.py b/deepmd/descriptor/se_a_type.py new file mode 100644 index 0000000000..31259b1043 --- /dev/null +++ b/deepmd/descriptor/se_a_type.py @@ -0,0 +1,807 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.utils.network import one_layer +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from deepmd.utils.network import embedding_net_type, embedding_net,share_embedding_network_oneside,share_embedding_network_twoside +from .se_a import DescrptSeA + +import math +from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, get_np_precision +from deepmd.utils.argcheck import list_to_doc +from deepmd.utils.tabulate import DeepTabulate +from deepmd.utils.type_embed import Type_embed_net + +class DescrptSeAEbd_type (): + def __init__ (self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24,48,96], + axis_neuron: int = 8, + resnet_dt: bool = False, + trainable: bool = True, + seed: int = 1, + type_one_side: bool = True, + type_filter: List[int]=[], + numb_aparam : int = 0, + set_davg_zero: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + exclude_types: List[int] = [], + ) -> None: + """ + Constructor + + Parameters + ---------- + rcut + The cut-off radius + rcut_smth + From where the environment matrix should be smoothed + sel : list[str] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + axis_neuron + Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix) + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + trainable + If the weights of embedding net are trainable. + seed + Random seed for initializing the network parameters. + type_one_side + Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets + type_nchanl + Number of channels for type representation + type_nlayer + Number of hidden layers for the type embedding net (skip connected). + numb_aparam + Number of atomic parameters. If >0 it will be embedded with atom types. + set_davg_zero + Set the shift of embedding net input to zero. + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + """ + # args = ClassArg()\ + # .add('type_nchanl', int, default = 4) \ + # .add('type_nlayer', int, default = 2) \ + # .add('type_one_side', bool, default = True) \ + # .add('numb_aparam', int, default = 0) + # class_data = args.parse(jdata) + self.sel_a = sel + self.rcut_r = rcut + self.rcut_r_smth = rcut_smth + self.filter_neuron = neuron + self.n_axis_neuron = axis_neuron + self.filter_resnet_dt = resnet_dt + self.seed = seed + self.trainable = trainable + self.filter_activation_fn = get_activation_func(activation_function) + self.filter_precision = get_precision(precision) + self.filter_np_precision = get_np_precision(precision) + self.exclude_types = set() + for tt in exclude_types: + assert(len(tt) == 2) + self.exclude_types.add((tt[0], tt[1])) + self.exclude_types.add((tt[1], tt[0])) + self.set_davg_zero = set_davg_zero + self.type_one_side = type_one_side + if self.type_one_side and len(exclude_types) != 0: + raise RuntimeError('"type_one_side" is not compatible with "exclude_types"') + + # descrpt config + self.sel_r = [ 0 for ii in range(len(self.sel_a)) ] + self.ntypes = len(self.sel_a) + assert(self.ntypes == len(self.sel_r)) + self.rcut_a = -1 + # numb of neighbors and numb of descrptors + self.nnei_a = np.cumsum(self.sel_a)[-1] + self.nnei_r = np.cumsum(self.sel_r)[-1] + self.nnei = self.nnei_a + self.nnei_r + self.ndescrpt_a = self.nnei_a * 4 + self.ndescrpt_r = self.nnei_r * 1 + self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r + self.useBN = False + self.dstd = None + self.davg = None + self.compress = False + self.place_holders = {} + avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + sub_graph = tf.Graph() + with sub_graph.as_default(): + name_pfx = 'd_sea_ebd_' + for ii in ['coord', 'box']: + self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii) + self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type') + self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms') + self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh') + self.stat_descrpt, descrpt_deriv, rij, nlist \ + = op_module.prod_env_mat_a(self.place_holders['coord'], + self.place_holders['type'], + self.place_holders['natoms_vec'], + self.place_holders['box'], + self.place_holders['default_mesh'], + tf.constant(avg_zero), + tf.constant(std_ones), + rcut_a = self.rcut_a, + rcut_r = self.rcut_r, + rcut_r_smth = self.rcut_r_smth, + sel_a = self.sel_a, + sel_r = self.sel_r) + self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config) + self.numb_aparam = numb_aparam + self.type_filter = type_filter + self.type_embed_net = Type_embed_net(self.ntypes,type_filter,resnet_dt,seed, activation_function,precision) + + if self.numb_aparam > 0: + add_data_requirement('aparam', 3, atomic=True, must=True, high_prec=False) + + def get_rcut (self) -> float: + """ + Returns the cut-off radius + """ + return self.rcut_r + + def get_ntypes (self) -> int: + """ + Returns the number of atom types + """ + return self.ntypes + + def get_dim_out (self) -> int: + """ + Returns the output dimension of this descriptor + """ + return self.filter_neuron[-1] * self.n_axis_neuron + + def get_dim_rot_mat_1 (self) -> int: + """ + Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3 + """ + return self.filter_neuron[-1] + + def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: + """ + Returns + ------- + nlist + Neighbor list + rij + The relative distance between the neighbor and the center atom. + sel_a + The number of neighbors with full information + sel_r + The number of neighbors with only radial information + """ + return self.nlist, self.rij, self.sel_a, self.sel_r + + def compute_input_stats (self, + data_coord : list, + data_box : list, + data_atype : list, + natoms_vec : list, + mesh : list, + input_dict : dict + ) -> None : + """ + Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + + Parameters + ---------- + data_coord + The coordinates. Can be generated by deepmd.model.make_stat_input + data_box + The box. Can be generated by deepmd.model.make_stat_input + data_atype + The atom types. Can be generated by deepmd.model.make_stat_input + natoms_vec + The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input + mesh + The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input + input_dict + Dictionary for additional input + """ + all_davg = [] + all_dstd = [] + if True: + sumr = [] + suma = [] + sumn = [] + sumr2 = [] + suma2 = [] + for cc,bb,tt,nn,mm in zip(data_coord,data_box,data_atype,natoms_vec,mesh) : + sysr,sysr2,sysa,sysa2,sysn \ + = self._compute_dstats_sys_smth(cc,bb,tt,nn,mm) + sumr.append(sysr) + suma.append(sysa) + sumn.append(sysn) + sumr2.append(sysr2) + suma2.append(sysa2) + sumr = np.sum(sumr, axis = 0) + suma = np.sum(suma, axis = 0) + sumn = np.sum(sumn, axis = 0) + sumr2 = np.sum(sumr2, axis = 0) + suma2 = np.sum(suma2, axis = 0) + for type_i in range(self.ntypes) : + davgunit = [sumr[type_i]/sumn[type_i], 0, 0, 0] + dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]) + ] + davg = np.tile(davgunit, self.ndescrpt // 4) + dstd = np.tile(dstdunit, self.ndescrpt // 4) + all_davg.append(davg) + all_dstd.append(dstd) + + if not self.set_davg_zero: + self.davg = np.array(all_davg) + self.dstd = np.array(all_dstd) + + def _compute_dstats_sys_smth (self, + data_coord, + data_box, + data_atype, + natoms_vec, + mesh) : + dd_all \ + = self.sub_sess.run(self.stat_descrpt, + feed_dict = { + self.place_holders['coord']: data_coord, + self.place_holders['type']: data_atype, + self.place_holders['natoms_vec']: natoms_vec, + self.place_holders['box']: data_box, + self.place_holders['default_mesh']: mesh, + }) + natoms = natoms_vec + dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]]) + start_index = 0 + sysr = [] + sysa = [] + sysn = [] + sysr2 = [] + sysa2 = [] + for type_i in range(self.ntypes): + end_index = start_index + self.ndescrpt * natoms[2+type_i] + dd = dd_all[:, start_index:end_index] + dd = np.reshape(dd, [-1, self.ndescrpt]) + start_index = end_index + # compute + dd = np.reshape (dd, [-1, 4]) + ddr = dd[:,:1] + dda = dd[:,1:] + sumr = np.sum(ddr) + suma = np.sum(dda) / 3. + sumn = dd.shape[0] + sumr2 = np.sum(np.multiply(ddr, ddr)) + suma2 = np.sum(np.multiply(dda, dda)) / 3. + sysr.append(sumr) + sysa.append(suma) + sysn.append(sumn) + sysr2.append(sumr2) + sysa2.append(suma2) + return sysr, sysr2, sysa, sysa2, sysn + + + def _compute_std (self,sumv2, sumv, sumn) : + val = np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn)) + if np.abs(val) < 1e-2: + val = 1e-2 + return val + def enable_compression(self, + min_nbor_dist : float, + model_file : str = 'frozon_model.pb', + table_extrapolate : float = 5, + table_stride_1 : float = 0.01, + table_stride_2 : float = 0.1, + check_frequency : int = -1 + ) -> None: + """ + Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data. + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + model_file + The original frozen model, which will be compressed by the program + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + self.compress = True + self.model_file = model_file + self.table_config = [table_extrapolate, table_stride_1, table_stride_2, check_frequency] + self.table = DeepTabulate(self.model_file, self.type_one_side) + self.lower, self.upper \ + = self.table.build(min_nbor_dist, + table_extrapolate, + table_stride_1, + table_stride_2) + + def prod_force_virial(self, + atom_ener : tf.Tensor, + natoms : tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial + + Parameters + ---------- + atom_ener + The atomic energy + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: number of local atoms + natoms[1]: total number of atoms held by this processor + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms + Return + ------ + force + The force on atoms + virial + The total virial + atom_virial + The atomic virial + """ + [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape) + tf.summary.histogram('net_derivative', net_deriv) + net_deriv_reshape = tf.reshape (net_deriv, [-1, natoms[0] * self.ndescrpt]) + force \ + = op_module.prod_force_se_a (net_deriv_reshape, + self.descrpt_deriv, + self.nlist, + natoms, + n_a_sel = self.nnei_a, + n_r_sel = self.nnei_r) + virial, atom_virial \ + = op_module.prod_virial_se_a (net_deriv_reshape, + self.descrpt_deriv, + self.rij, + self.nlist, + natoms, + n_a_sel = self.nnei_a, + n_r_sel = self.nnei_r) + tf.summary.histogram('force', force) + tf.summary.histogram('virial', virial) + tf.summary.histogram('atom_virial', atom_virial) + + return force, virial, atom_virial + + def build (self, + coord_ : tf.Tensor, + atype_ : tf.Tensor, + natoms : tf.Tensor, + box_ : tf.Tensor, + mesh : tf.Tensor, + input_dict : dict, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor + + Parameters + ---------- + coord_ + The coordinate of atoms + atype_ + The type of atoms + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: number of local atoms + natoms[1]: total number of atoms held by this processor + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms + mesh + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict + Dictionary for additional inputs + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + descriptor + The output descriptor + """ + nei_type = np.array([]) + for ii in range(self.ntypes): + nei_type = np.append(nei_type, ii * np.ones(self.sel_a[ii])) # like a mask + self.nei_type = tf.get_variable('t_nei_type', + [self.nnei], + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(nei_type)) + + davg = self.davg + dstd = self.dstd + #self.type_filter[-1] = self.filter_neuron[0] + with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) : + if davg is None: + davg = np.zeros([self.ntypes, self.ndescrpt]) + if dstd is None: + dstd = np.ones ([self.ntypes, self.ndescrpt]) + t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), + name = 'rcut', + dtype = GLOBAL_TF_FLOAT_PRECISION) + t_ntypes = tf.constant(self.ntypes, + name = 'ntypes', + dtype = tf.int32) + t_ndescrpt = tf.constant(self.ndescrpt, + name = 'ndescrpt', + dtype = tf.int32) + t_sel = tf.constant(self.sel_a, + name = 'sel', + dtype = tf.int32) + self.t_avg = tf.get_variable('t_avg', + davg.shape, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(davg)) + self.t_std = tf.get_variable('t_std', + dstd.shape, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(dstd)) + + coord = tf.reshape (coord_, [-1, natoms[1] * 3]) + box = tf.reshape (box_, [-1, 9]) + atype = tf.reshape (atype_, [-1, natoms[1]]) + + self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ + = op_module.prod_env_mat_a (coord, + atype, + natoms, + box, + mesh, + self.t_avg, + self.t_std, + rcut_a = self.rcut_a, + rcut_r = self.rcut_r, + rcut_r_smth = self.rcut_r_smth, + sel_a = self.sel_a, + sel_r = self.sel_r) + # only used when tensorboard was set as true + tf.summary.histogram('descrpt', self.descrpt) + tf.summary.histogram('rij', self.rij) + tf.summary.histogram('nlist', self.nlist) + + self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) + self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat') + self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv') + self.rij = tf.identity(self.rij, name = 'o_rij') + self.nlist = tf.identity(self.nlist, name = 'o_nlist') + + self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, + atype, + natoms, + input_dict, + suffix = suffix, + reuse = reuse, + trainable = self.trainable) + + # only used when tensorboard was set as true + tf.summary.histogram('embedding_net_output', self.dout) + return self.dout + + + def _embedding_net(self, + inputs, + natoms, + filter_neuron, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + seed=None, + trainable = True): + ''' + inputs: nf x na x (nei x 4) + outputs: nf x na x nei x output_size + ''' + # natom x (nei x 4) + inputs = tf.reshape(inputs, [-1, self.ndescrpt]) + shape = inputs.get_shape().as_list() + outputs_size = [1] + filter_neuron + with tf.variable_scope(name, reuse=reuse): + xyz_scatter_total = [] + # with natom x (nei x 4) + inputs_i = inputs + shape_i = inputs_i.get_shape().as_list() + # with (natom x nei) x 4 + inputs_reshape = tf.reshape(inputs_i, [-1, 4]) # Ri + # with (natom x nei) x 1 + xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0,0],[-1,1]),[-1,1]) # the col of sij + # with (natom x nei) x out_size + xyz_scatter = embedding_net(xyz_scatter, + self.filter_neuron, + self.filter_precision, + activation_fn = activation_fn, + resnet_dt = self.filter_resnet_dt, + stddev = stddev, + bavg = bavg, + seed = seed, + trainable = trainable) + # natom x nei x out_size + xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1])) + xyz_scatter_total.append(xyz_scatter) + # natom x nei x outputs_size + xyz_scatter = tf.concat(xyz_scatter_total, axis=1) + # nf x natom x nei x outputs_size + xyz_scatter = tf.reshape(xyz_scatter, [tf.shape(inputs)[0], natoms[0], self.nnei, outputs_size[-1]]) + return xyz_scatter + + def _share_embedding_net_twoside(self, + inputs, + natoms, + atype, + nframes, + filter_neuron, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + seed=None, + trainable = True): + ''' + inputs: nf x na x (nei x 4) + outputs: nf x na x nei x output_size + ''' + # natom x (nei x 4) + #init_shape = inputs.get_shape().as_list() + #nf = init_shape[0] + + inputs = tf.reshape(inputs, [-1, self.ndescrpt]) + shape = inputs.get_shape().as_list() + outputs_size = [1] + filter_neuron + with tf.variable_scope(name, reuse=reuse): + xyz_scatter_total = [] + # with natom x (nei x 4) + inputs_i = inputs + shape_i = inputs_i.get_shape().as_list() + # with (natom x nei) x 4 + inputs_reshape = tf.reshape(inputs_i, [-1, 4]) # Ri + # with (natom x nei) x 1 + xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0,0],[-1,1]),[-1,1]) # the col of sij + + # with [ntypes, nchanl] + + #self.nei_type = tf.reshape(self.nei_type,[-1,1]) + + nei_embed = self.type_embed_net._type_embed( self.nei_type, + reuse = reuse, + trainable = True, + suffix = '') #nnei*nchnl + atm_embed = self.type_embed_net._type_embed( atype, + reuse = True, + trainable = True, + suffix = '') #(nf*natom)*nchnl + nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) + nei_embed = tf.reshape(nei_embed,[-1,self.type_filter[-1]]) + + atm_embed = tf.tile(atm_embed,(1,self.nnei)) + atm_embed = tf.reshape(atm_embed,[-1,self.type_filter[-1]]) + filter_tmp = [1+2*self.type_filter[-1]]+self.filter_neuron + + xyz_scatter = share_embedding_network_twoside(xyz_scatter, + atype, + self.nei_type, + nei_embed, + atm_embed, + filter_tmp, + self.filter_precision, + activation_fn = activation_fn, + resnet_dt = self.filter_resnet_dt, + stddev = stddev, + bavg = bavg, + seed = seed, + trainable = trainable) + # with (natom x nei) x out_size + # natom x nei x out_size + + xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1]) ) + xyz_scatter_total.append(xyz_scatter) + # natom x nei x outputs_size + xyz_scatter = tf.concat(xyz_scatter_total, axis=1) + # nf x natom x nei x outputs_size + xyz_scatter = tf.reshape(xyz_scatter, [tf.shape(inputs)[0], natoms[0], self.nnei, outputs_size[-1]]) + + return xyz_scatter + + def _share_embedding_net_oneside(self, + inputs, + natoms, + atype, + nframes, + filter_neuron, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + seed=None, + trainable = True): + ''' + inputs: nf x na x (nei x 4) + outputs: nf x na x nei x output_size + ''' + # natom x (nei x 4) + #init_shape = inputs.get_shape().as_list() + #nf = init_shape[0] + + inputs = tf.reshape(inputs, [-1, self.ndescrpt]) + shape = inputs.get_shape().as_list() + outputs_size = [1] + filter_neuron + with tf.variable_scope(name, reuse=reuse): + xyz_scatter_total = [] + # with natom x (nei x 4) + inputs_i = inputs + shape_i = inputs_i.get_shape().as_list() + # with (natom x nei) x 4 + inputs_reshape = tf.reshape(inputs_i, [-1, 4]) # Ri + # with (natom x nei) x 1 + xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0,0],[-1,1]),[-1,1]) # the col of sij + # with [ntypes, nchanl] + + + #self.nei_type = tf.reshape(self.nei_type,[-1,1]) + + nei_embed = self.type_embed_net._type_embed( self.nei_type, + reuse = reuse, + trainable = True, + suffix = '') #nnei*nchnl + + + nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) + nei_embed = tf.reshape(nei_embed,[-1,self.type_filter[-1]]) + + filter_tmp = [1+self.type_filter[-1]]+self.filter_neuron + + xyz_scatter = share_embedding_network_oneside(xyz_scatter, + atype, + self.nei_type, + nei_embed, + filter_tmp, + self.filter_precision, + activation_fn = activation_fn, + resnet_dt = self.filter_resnet_dt, + stddev = stddev, + bavg = bavg, + seed = seed, + trainable = trainable) + # with (natom x nei) x out_size + # natom x nei x out_size + + xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1]) ) + xyz_scatter_total.append(xyz_scatter) + # natom x nei x outputs_size + xyz_scatter = tf.concat(xyz_scatter_total, axis=1) + # nf x natom x nei x outputs_size + xyz_scatter = tf.reshape(xyz_scatter, [tf.shape(inputs)[0], natoms[0], self.nnei, outputs_size[-1]]) + + return xyz_scatter + + + + + + + def _pass_filter(self, + inputs, + atype, + natoms, + input_dict, + reuse = None, + suffix = '', + trainable = True) : + # nf x na x ndescrpt + # nf x na x (nnei x 4) + + inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]) # total input + layer, qmat = self._share_filter(tf.cast(inputs, self.filter_precision), + atype, + natoms, + input_dict, + name='filter_type_all'+suffix, + reuse=reuse, + seed = self.seed, + trainable = trainable, + activation_fn = self.filter_activation_fn) + output = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) + output_qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3]) + return output, output_qmat + + + + + def _share_filter(self, + inputs, + atype, + natoms, + input_dict, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + seed=None, + trainable = True): + outputs_size = self.filter_neuron[-1] + outputs_size_2 = self.n_axis_neuron + # nf x natom x (nei x 4) + nframes = tf.shape(inputs)[0] + + shape = tf.reshape(inputs, [-1, self.ndescrpt]).get_shape().as_list() #[natom, ndescrpt] + + # nf x natom x nei x outputs_size + if self.type_one_side: + mat_g = self._share_embedding_net_oneside(inputs, + natoms, + atype, + nframes, + self.filter_neuron, + activation_fn = activation_fn, + stddev = stddev, + bavg = bavg, + name = name, + reuse = reuse, + seed = seed, + trainable = trainable) + else: + mat_g = self._share_embedding_net_twoside(inputs, + natoms, + atype, + nframes, + self.filter_neuron, + activation_fn = activation_fn, + stddev = stddev, + bavg = bavg, + name = name, + reuse = reuse, + seed = seed, + trainable = trainable) + # (nf x natom) x nei x outputs_size + mat_g = tf.reshape(mat_g, [nframes*natoms[0], self.nnei, outputs_size]) + # natom x nei x 4 + inputs_reshape = tf.reshape(inputs, [-1, shape[1]//4, 4]) + + # natom x 4 x outputs_size + xyz_scatter_1 = tf.matmul(inputs_reshape, mat_g, transpose_a = True) + + xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape[1]) + # natom x 4 x outputs_size_2 + xyz_scatter_2 = tf.slice(xyz_scatter_1, [0,0,0],[-1,-1,outputs_size_2]) + # # natom x 3 x outputs_size_2 + # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -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]) + # 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]) + + return result, qmat + diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index 5d60ffcc09..34353cc03b 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -26,7 +26,8 @@ def __init__ (self, seed : int = 1, atom_ener : List[float] = [], activation_function : str = 'tanh', - precision : str = 'default' + precision : str = 'default', + share_fitting : bool = False ) -> None: """ Constructor @@ -86,6 +87,7 @@ def __init__ (self, self.fitting_activation_fn = get_activation_func(activation_function) self.fitting_precision = get_precision(precision) self.trainable = trainable + self.share_fitting = share_fitting if self.trainable is None: self.trainable = [True for ii in range(len(self.n_neuron) + 1)] if type(self.trainable) is bool: @@ -366,3 +368,127 @@ def build (self, + def build_share (self, + inputs : tf.Tensor, + atype: tf.Tensor, + type_embedding: tf.Tensor, + natoms : tf.Tensor, + input_dict : dict = {}, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for fitting net + + Parameters + ---------- + inputs + The input descriptor + input_dict + Additional dict for inputs. + if numb_fparam > 0, should have input_dict['fparam'] + if numb_aparam > 0, should have input_dict['aparam'] + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: number of local atoms + natoms[1]: total number of atoms held by this processor + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Return + ------ + ener + The system energy + """ + bias_atom_e = self.bias_atom_e + if self.numb_fparam > 0 and ( self.fparam_avg is None or self.fparam_inv_std is None ): + raise RuntimeError('No data stat result. one should do data statisitic, before build') + if self.numb_aparam > 0 and ( self.aparam_avg is None or self.aparam_inv_std is None ): + raise RuntimeError('No data stat result. one should do data statisitic, before build') + + with tf.variable_scope('fitting_attr' + suffix, reuse = reuse) : + t_dfparam = tf.constant(self.numb_fparam, + name = 'dfparam', + dtype = tf.int32) + t_daparam = tf.constant(self.numb_aparam, + name = 'daparam', + dtype = tf.int32) + if self.numb_fparam > 0: + t_fparam_avg = tf.get_variable('t_fparam_avg', + self.numb_fparam, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(self.fparam_avg)) + t_fparam_istd = tf.get_variable('t_fparam_istd', + self.numb_fparam, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(self.fparam_inv_std)) + if self.numb_aparam > 0: + t_aparam_avg = tf.get_variable('t_aparam_avg', + self.numb_aparam, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(self.aparam_avg)) + t_aparam_istd = tf.get_variable('t_aparam_istd', + self.numb_aparam, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(self.aparam_inv_std)) + + atype_shape = atype.get_shape().as_list() + type_shape = type_embedding.get_shape().as_list() + atm_embed = tf.nn.embedding_lookup(type_embedding,atype) + atm_embed = tf.reshape(atm_embed,[-1,type_shape[1]]) + inputs = tf.concat([tf.reshape(inputs,[-1,self.dim_descrpt]),atm_embed],axis=1) + start_index = 0 + inputs = tf.cast(tf.reshape(inputs, [-1, (self.dim_descrpt+type_shape[1]) * natoms[0]]), self.fitting_precision) + + if bias_atom_e is not None : + assert(len(bias_atom_e) == self.ntypes) + + if self.numb_fparam > 0 : + fparam = input_dict['fparam'] + fparam = tf.reshape(fparam, [-1, self.numb_fparam]) + fparam = (fparam - t_fparam_avg) * t_fparam_istd + if self.numb_aparam > 0 : + aparam = input_dict['aparam'] + aparam = tf.reshape(aparam, [-1, self.numb_aparam]) + aparam = (aparam - t_aparam_avg) * t_aparam_istd + aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]]) + + layer = tf.reshape(inputs, [-1, self.dim_descrpt+type_shape[1]]) + type_bias_ae=0 + if self.numb_fparam > 0 : + ext_fparam = tf.tile(fparam, [1, natoms[0]]) + ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam]) + ext_fparam = tf.cast(ext_fparam,self.fitting_precision) + layer = tf.concat([layer, ext_fparam], axis = 1) + if self.numb_aparam > 0 : + ext_aparam = tf.reshape(ext_aparam, [-1, self.numb_aparam]) + ext_aparam = tf.cast(ext_aparam,self.fitting_precision) + layer = tf.concat([layer, ext_aparam], axis = 1) + + for ii in range(0,len(self.n_neuron)) : + if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : + layer+= one_layer(layer, self.n_neuron[ii], name='fitting_layer_'+str(ii)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii]) + else : + layer = one_layer(layer, self.n_neuron[ii], name='fitting_layer_'+str(ii)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, trainable = self.trainable[ii]) + final_layer = one_layer(layer, 1, activation_fn = None, bavg = type_bias_ae, name='fitting_final_layer_'+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, trainable = self.trainable[-1]) + + + outs = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[0]]) + + if self.tot_ener_zero: + force_tot_ener = 0.0 + outs = tf.reshape(outs, [-1, natoms[0]]) + outs_mean = tf.reshape(tf.reduce_mean(outs, axis = 1), [-1, 1]) + outs_mean = outs_mean - tf.ones_like(outs_mean, dtype = GLOBAL_TF_FLOAT_PRECISION) * (force_tot_ener/global_cvt_2_tf_float(natoms[0])) + outs = outs - outs_mean + outs = tf.reshape(outs, [-1]) + + tf.summary.histogram('fitting_net_output', outs) + return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) \ No newline at end of file diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py index 3cc3d4bd8b..c5fedbfff1 100644 --- a/deepmd/model/ener.py +++ b/deepmd/model/ener.py @@ -7,6 +7,7 @@ from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION from deepmd.env import op_module from .model_stat import make_stat_input, merge_sys_stat +from deepmd.descriptor import DescrptSeAEbd_type class EnerModel() : model_type = 'ener' @@ -138,6 +139,7 @@ def build (self, coord = tf.reshape (coord_, [-1, natoms[1] * 3]) atype = tf.reshape (atype_, [-1, natoms[1]]) + dout \ = self.descrpt.build(coord_, atype_, @@ -145,21 +147,31 @@ def build (self, box, mesh, input_dict, - suffix = suffix, - reuse = reuse) + reuse = reuse, + suffix = suffix) dout = tf.identity(dout, name='o_descriptor') - + if isinstance(self.descrpt,DescrptSeAEbd_type): + type_embedding = self.descrpt.type_embed_net.fetch_type_embedding() + type_embedding = tf.identity(type_embedding,name ='t_embed') 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, + + if self.fitting.share_fitting: + atom_ener = self.fitting.build_share(dout, + atype_, + type_embedding, natoms, input_dict, reuse = reuse, suffix = suffix) - + else: + atom_ener = self.fitting.build (dout, + natoms, + input_dict, + reuse = reuse, + suffix = suffix) if self.srtab is not None : sw_lambda, sw_deriv \ = op_module.soft_min_switch(atype, diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index 62934b0faa..2eebec1e8c 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -12,7 +12,7 @@ from deepmd.descriptor import DescrptLocFrame from deepmd.descriptor import DescrptSeA from deepmd.descriptor import DescrptSeT -from deepmd.descriptor import DescrptSeAEbd +from deepmd.descriptor import DescrptSeAEbd,DescrptSeAEbd_type from deepmd.descriptor import DescrptSeAEf from deepmd.descriptor import DescrptSeR from deepmd.descriptor import DescrptSeAR @@ -61,8 +61,10 @@ def _generate_descrpt_from_param_dict(descrpt_param): descrpt = DescrptSeR(**descrpt_param) elif descrpt_type == 'se_e3' or descrpt_type == 'se_at' or descrpt_type == 'se_a_3be' : descrpt = DescrptSeT(**descrpt_param) - elif descrpt_type == 'se_a_tpe' or descrpt_type == 'se_a_ebd' : + elif descrpt_type == 'se_a_tpe' : descrpt = DescrptSeAEbd(**descrpt_param) + elif descrpt_type == 'se_a_ebd': + descrpt = DescrptSeAEbd_type(**descrpt_param) elif descrpt_type == 'se_a_ef' : descrpt = DescrptSeAEf(**descrpt_param) elif descrpt_type == 'se_ar' : @@ -474,6 +476,8 @@ def train (self, train_data, valid_data=None) : tb_valid_writer = None train_time = 0 + varvar2 = self.sess.graph.get_tensor_by_name('t_embed:0') + print(self.sess.run(varvar2)) while cur_batch < stop_batch : # first round validation: @@ -511,6 +515,8 @@ def train (self, train_data, valid_data=None) : log.info("batch %7d training time %.2f s, testing time %.2f s" % (cur_batch, train_time, test_time)) train_time = 0 + varvar2 = self.sess.graph.get_tensor_by_name('t_embed:0') + print(self.sess.run(varvar2)) if self.save_freq > 0 and cur_batch % self.save_freq == 0 and self.run_opt.is_chief : if self.saver is not None : self.saver.save (self.sess, os.getcwd() + "/" + self.save_ckpt) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a29fd3b7f2..8227520bf5 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -109,6 +109,38 @@ def descrpt_se_a_tpe_args(): Argument("numb_aparam", int, optional = True, default = 0, doc = doc_numb_aparam) ] +def descrpt_se_a_ebd_args(): + doc_sel = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.' + doc_rcut = 'The cut-off radius.' + doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`' + doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.' + doc_axis_neuron = 'Size of the submatrix of G (embedding matrix).' + doc_type_filter = 'Size of type embed net' + doc_activation_function = f'The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}' + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_type_one_side = 'Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets' + doc_precision = f'The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' + doc_trainable = 'If the parameters in the embedding net is trainable' + doc_seed = 'Random seed for parameter initialization' + doc_exclude_types = 'The Excluded types' + doc_set_davg_zero = 'Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used' + + return [ + Argument("sel", list, optional = False, doc = doc_sel), + Argument("rcut", float, optional = True, default = 6.0, doc = doc_rcut), + Argument("rcut_smth", float, optional = True, default = 0.5, doc = doc_rcut_smth), + Argument("neuron", list, optional = True, default = [10,20,40], doc = doc_neuron), + Argument("axis_neuron", int, optional = True, default = 4, doc = doc_axis_neuron), + Argument("activation_function", str, optional = True, default = 'tanh', doc = doc_activation_function), + Argument("resnet_dt", bool, optional = True, default = False, doc = doc_resnet_dt), + Argument("type_one_side", bool, optional = True, default = False, doc = doc_type_one_side), + Argument("type_filter", list, optional = True, default = [5,10,10], doc = doc_type_filter), + Argument("precision", str, optional = True, default = "float64", doc = doc_precision), + Argument("trainable", bool, optional = True, default = True, doc = doc_trainable), + Argument("seed", [int,None], optional = True, doc = doc_seed), + Argument("exclude_types", list, optional = True, default = [], doc = doc_exclude_types), + Argument("set_davg_zero", bool, optional = True, default = False, doc = doc_set_davg_zero) + ] def descrpt_se_r_args(): doc_sel = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.' @@ -166,6 +198,7 @@ def descrpt_variant_type_args(): link_se_e2_r = make_link('se_e2_r', 'model/descriptor[se_e2_r]') link_se_e3 = make_link('se_e3', 'model/descriptor[se_e3]') link_se_a_tpe = make_link('se_a_tpe', 'model/descriptor[se_a_tpe]') + link_se_a_ebd = make_link('se_a_ebd', 'model/descriptor[se_a_ebd]') link_hybrid = make_link('hybrid', 'model/descriptor[hybrid]') doc_descrpt_type = f'The type of the descritpor. See explanation below. \n\n\ - `loc_frame`: Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.\n\n\ @@ -173,6 +206,7 @@ def descrpt_variant_type_args(): - `se_e2_r`: Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor.\n\n\ - `se_e3`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Three-body embedding will be used by this descriptor.\n\n\ - `se_a_tpe`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Type embedding will be used by this descriptor.\n\n\ +- `se_a_ebd`: Type embedding will be used by this descriptor.\n\n\ - `hybrid`: Concatenate of a list of descriptors as a new descriptor.' return Variant("type", [ @@ -180,7 +214,8 @@ def descrpt_variant_type_args(): Argument("se_e2_a", dict, descrpt_se_a_args(), alias = ['se_a']), Argument("se_e2_r", dict, descrpt_se_r_args(), alias = ['se_r']), Argument("se_e3", dict, descrpt_se_t_args(), alias = ['se_at', 'se_a_3be', 'se_t']), - Argument("se_a_tpe", dict, descrpt_se_a_tpe_args(), alias = ['se_a_ebd']), + Argument("se_a_tpe", dict, descrpt_se_a_tpe_args()), + Argument("se_a_ebd", dict, descrpt_se_a_ebd_args(), alias = ['se_a_ebd_type']), Argument("hybrid", dict, descrpt_hybrid_args()), ], doc = doc_descrpt_type) @@ -193,6 +228,7 @@ def fitting_ener(): doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}' doc_precision = f'The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_share_fitting = 'Whether to share a single fitting net' doc_trainable = 'Whether the parameters in the fitting net are trainable. This option can be\n\n\ - bool: True if all parameters of the fitting net are trainable, False otherwise.\n\n\ - list of bool: Specifies if each layer is trainable. Since the fitting net is composed by hidden layers followed by a output layer, the length of tihs list should be equal to len(`neuron`)+1.' @@ -207,6 +243,7 @@ def fitting_ener(): Argument("activation_function", str, optional = True, default = 'tanh', doc = doc_activation_function), Argument("precision", str, optional = True, default = 'float64', doc = doc_precision), Argument("resnet_dt", bool, optional = True, default = True, doc = doc_resnet_dt), + Argument("share_fitting", bool, optional = True, default = False, doc = doc_share_fitting), Argument("trainable", [list,bool], optional = True, default = True, doc = doc_trainable), Argument("rcond", float, optional = True, default = 1e-3, doc = doc_rcond), Argument("seed", [int,None], optional = True, doc = doc_seed), diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py index ab32308bdd..e84062756e 100644 --- a/deepmd/utils/network.py +++ b/deepmd/utils/network.py @@ -92,6 +92,297 @@ def embedding_net(xx, """ outputs_size = [1] + network_size + for ii in range(1, len(outputs_size)): + w = tf.get_variable('matrix_'+str(ii)+name_suffix, + [outputs_size[ii - 1], outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), seed = seed), + trainable = trainable) + variable_summaries(w, 'matrix_'+str(ii)+name_suffix) + + b = tf.get_variable('bias_'+str(ii)+name_suffix, + [1, outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), + trainable = trainable) + variable_summaries(b, 'bias_'+str(ii)+name_suffix) + + hidden = tf.reshape(activation_fn(tf.matmul(xx, w) + b), [-1, outputs_size[ii]]) + if resnet_dt : + idt = tf.get_variable('idt_'+str(ii)+name_suffix, + [1, outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=0.001, mean = 1.0, seed = seed), + trainable = trainable) + variable_summaries(idt, 'idt_'+str(ii)+name_suffix) + + if outputs_size[ii] == outputs_size[ii-1]: + if resnet_dt : + xx += hidden * idt + else : + xx += hidden + elif outputs_size[ii] == outputs_size[ii-1] * 2: + if resnet_dt : + xx = tf.concat([xx,xx], 1) + hidden * idt + else : + xx = tf.concat([xx,xx], 1) + hidden + else: + xx = hidden + + return xx + +def embedding_net_type(xx, + network_size, + precision, + activation_fn = tf.nn.tanh, + resnet_dt = False, + name_suffix = '', + stddev = 1.0, + bavg = 0.0, + seed = None, + trainable = True): + """ + Parameters + ---------- + xx : Tensor + Input tensor of shape [-1,1] + network_size: list of int + Size of the embedding network. For example [16,32,64] + precision: + Precision of network weights. For example, tf.float64 + activation_fn: + Activation function + resnet_dt: boolean + Using time-step in the ResNet construction + name_suffix: str + The name suffix append to each variable. + stddev: float + Standard deviation of initializing network parameters + bavg: float + Mean of network intial bias + seed: int + Random seed for initializing network parameters + trainable: boolean + If the netowk is trainable + """ + outputs_size = network_size + + for ii in range(1, len(outputs_size)): + w = tf.get_variable('matrix_'+str(ii)+name_suffix, + [outputs_size[ii - 1], outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), seed = seed), + trainable = trainable) + variable_summaries(w, 'matrix_'+str(ii)+name_suffix) + + b = tf.get_variable('bias_'+str(ii)+name_suffix, + [1, outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), + trainable = trainable) + variable_summaries(b, 'bias_'+str(ii)+name_suffix) + + hidden = tf.reshape(activation_fn(tf.matmul(xx, w) + b), [-1, outputs_size[ii]]) + if resnet_dt : + idt = tf.get_variable('idt_'+str(ii)+name_suffix, + [1, outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=0.001, mean = 1.0, seed = seed), + trainable = trainable) + variable_summaries(idt, 'idt_'+str(ii)+name_suffix) + + if outputs_size[ii] == outputs_size[ii-1]: + if resnet_dt : + xx += hidden * idt + else : + xx += hidden + elif outputs_size[ii] == outputs_size[ii-1] * 2: + if resnet_dt : + xx = tf.concat([xx,xx], 1) + hidden * idt + else : + xx = tf.concat([xx,xx], 1) + hidden + else: + xx = hidden + + return xx + +def share_embedding_network_oneside(xx, + atype, + nei_type, + nei_embed, + network_size, + precision, + activation_fn = tf.nn.tanh, + resnet_dt = False, + name_suffix = '', + stddev = 1.0, + bavg = 0.0, + seed = None, + trainable = True): + """ + Parameters + ---------- + xx : Tensor + Input tensor of shape [-1,1] + atype: Tensor + atom type (length is natom) + nei_type: Tensor + The type mask of neighbors + network_size: list of int + Size of the embedding network. For example [16,32,64] + precision: + Precision of network weights. For example, tf.float64 + activation_fn: + Activation function + resnet_dt: boolean + Using time-step in the ResNet construction + name_suffix: str + The name suffix append to each variable. + stddev: float + Standard deviation of initializing network parameters + bavg: float + Mean of network intial bias + seed: int + Random seed for initializing network parameters + trainable: boolean + If the netowk is trainable + """ + + """ + outputs_size = [1] + network_size + w = tf.get_variable('matrix_'+str(1)+name_suffix, + [outputs_size[0], outputs_size[1]], + precision, + tf.random_normal_initializer(stddev=stddev/np.sqrt(outputs_size[1]+outputs_size[0]), seed = seed), + trainable = trainable) + variable_summaries(w, 'matrix_'+str(1)+name_suffix) + + b = tf.get_variable('bias_'+str(1)+name_suffix, + [1, outputs_size[1]], + precision, + tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), + trainable = trainable) + variable_summaries(b, 'bias_'+str(1)+name_suffix) + + hidden = tf.reshape(activation_fn(tf.matmul(xx, w) + b), [-1, outputs_size[1]]) # atom * nnei , output_size[1] + + xx = hidden + nei_embed + atm_embed + """ + + + outputs_size = network_size + xx = tf.concat([xx,nei_embed],1) + + + for ii in range(1, len(outputs_size)): + w = tf.get_variable('matrix_'+str(ii)+name_suffix, + [outputs_size[ii - 1], outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), seed = seed), + trainable = trainable) + variable_summaries(w, 'matrix_'+str(ii)+name_suffix) + + b = tf.get_variable('bias_'+str(ii)+name_suffix, + [1, outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), + trainable = trainable) + variable_summaries(b, 'bias_'+str(ii)+name_suffix) + + hidden = tf.reshape(activation_fn(tf.matmul(xx, w) + b), [-1, outputs_size[ii]]) + if resnet_dt : + idt = tf.get_variable('idt_'+str(ii)+name_suffix, + [1, outputs_size[ii]], + precision, + tf.random_normal_initializer(stddev=0.001, mean = 1.0, seed = seed), + trainable = trainable) + variable_summaries(idt, 'idt_'+str(ii)+name_suffix) + + if outputs_size[ii] == outputs_size[ii-1]: + if resnet_dt : + xx += hidden * idt + else : + xx += hidden + elif outputs_size[ii] == outputs_size[ii-1] * 2: + if resnet_dt : + xx = tf.concat([xx,xx], 1) + hidden * idt + else : + xx = tf.concat([xx,xx], 1) + hidden + else: + xx = hidden + + return xx + + +def share_embedding_network_twoside(xx, + atype, + nei_type, + nei_embed, + atm_embed, + network_size, + precision, + activation_fn = tf.nn.tanh, + resnet_dt = False, + name_suffix = '', + stddev = 1.0, + bavg = 0.0, + seed = None, + trainable = True): + """ + Parameters + ---------- + xx : Tensor + Input tensor of shape [-1,1] + atype: Tensor + atom type (length is natom) + nei_type: Tensor + The type mask of neighbors + network_size: list of int + Size of the embedding network. For example [16,32,64] + precision: + Precision of network weights. For example, tf.float64 + activation_fn: + Activation function + resnet_dt: boolean + Using time-step in the ResNet construction + name_suffix: str + The name suffix append to each variable. + stddev: float + Standard deviation of initializing network parameters + bavg: float + Mean of network intial bias + seed: int + Random seed for initializing network parameters + trainable: boolean + If the netowk is trainable + """ + + """ + outputs_size = [1] + network_size + w = tf.get_variable('matrix_'+str(1)+name_suffix, + [outputs_size[0], outputs_size[1]], + precision, + tf.random_normal_initializer(stddev=stddev/np.sqrt(outputs_size[1]+outputs_size[0]), seed = seed), + trainable = trainable) + variable_summaries(w, 'matrix_'+str(1)+name_suffix) + + b = tf.get_variable('bias_'+str(1)+name_suffix, + [1, outputs_size[1]], + precision, + tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), + trainable = trainable) + variable_summaries(b, 'bias_'+str(1)+name_suffix) + + hidden = tf.reshape(activation_fn(tf.matmul(xx, w) + b), [-1, outputs_size[1]]) # atom * nnei , output_size[1] + + xx = hidden + nei_embed + atm_embed + """ + + + outputs_size = network_size + xx = tf.concat([xx,nei_embed,atm_embed],1) + + for ii in range(1, len(outputs_size)): w = tf.get_variable('matrix_'+str(ii)+name_suffix, [outputs_size[ii - 1], outputs_size[ii]], diff --git a/deepmd/utils/type_embed.py b/deepmd/utils/type_embed.py new file mode 100644 index 0000000000..f6a8b942fe --- /dev/null +++ b/deepmd/utils/type_embed.py @@ -0,0 +1,66 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.utils.network import one_layer +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from deepmd.utils.network import embedding_net_type, embedding_net,share_embedding_network_oneside,share_embedding_network_twoside + +import math +from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, get_np_precision +from deepmd.utils.argcheck import list_to_doc +from deepmd.utils.tabulate import DeepTabulate + +class Type_embed_net(): + def __init__(self, + ntypes: int, + type_filter: List[int]=[], + resnet_dt: bool = False, + seed: int = 1, + activation_function: str = 'tanh', + precision: str = 'default', + )->None: + self.ntypes = ntypes + self.type_filter = type_filter + self.seed = seed + self.filter_resnet_dt = resnet_dt + self.filter_precision = get_precision(precision) + self.filter_np_precision = get_np_precision(precision) + self.filter_activation_fn = get_activation_func(activation_function) + + + def _type_embed(self, + atype, + reuse = None, + suffix = '', + trainable = True): + + ebd_type = tf.cast(tf.one_hot(tf.cast(atype,dtype=tf.int32),int(self.ntypes)), self.filter_precision) + ebd_type = tf.reshape(ebd_type, [-1, self.ntypes]) + name = 'type_embed_net_' + with tf.variable_scope(name, reuse=tf.AUTO_REUSE): + ebd_type = embedding_net_type(ebd_type, + [self.ntypes]+self.type_filter, + activation_fn = self.filter_activation_fn, + precision = self.filter_precision, + resnet_dt = self.filter_resnet_dt, + seed = self.seed, + trainable = trainable) + + ebd_type = tf.reshape(ebd_type, [-1, self.type_filter[-1]]) # nnei * type_filter[-1] + return ebd_type + + + def fetch_type_embedding(self) -> tf.Tensor: + _atom_type_ = [] + for ii in range(self.ntypes): + _atom_type_.append(ii) + _atom_type_ = tf.convert_to_tensor(_atom_type_,dtype = GLOBAL_TF_FLOAT_PRECISION) + type_embedding = self._type_embed( tf.one_hot(tf.cast(_atom_type_,dtype=tf.int32),int(self.ntypes)), + reuse = True, + suffix = '', + trainable = True) + return type_embedding \ No newline at end of file diff --git a/examples/water/se_a_type/input.json b/examples/water/se_a_type/input.json new file mode 100644 index 0000000000..454c98b5e7 --- /dev/null +++ b/examples/water/se_a_type/input.json @@ -0,0 +1,69 @@ +{ + "_comment": " model parameters", + "model": { + "type_map": ["O", "H"], + "descriptor" :{ + "type": "se_a_ebd", + "sel": [46, 92], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": false, + "axis_neuron": 16, + "type_filter": [5,10,10], + "type_one_side": false, + "seed": 1, + "_comment": " that's all" + }, + "fitting_net" : { + "neuron": [240, 240, 240], + "resnet_dt": true, + "seed": 1, + "share_fitting": true, + "_comment": " that's all" + }, + "_comment": " that's all" + }, + + "learning_rate" :{ + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment": "that's all" + }, + + "loss" :{ + "type": "ener", + "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, + "_comment": " that's all" + }, + + "training" : { + "training_data": { + "systems": ["../data/data_0/", "../data/data_1/", "../data/data_2/"], + "batch_size": "auto", + "_comment": "that's all" + }, + "validation_data":{ + "systems": ["../data/data_3"], + "batch_size": 1, + "numb_btch": 3, + "_comment": "that's all" + }, + "numb_steps": 1000000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 1000, + "_comment": "that's all" + }, + + "_comment": "that's all" +} + diff --git a/source/tests/e.out b/source/tests/e.out new file mode 100644 index 0000000000..4a758ca012 --- /dev/null +++ b/source/tests/e.out @@ -0,0 +1 @@ +4.826771866004193612e+01 diff --git a/source/tests/f.out b/source/tests/f.out new file mode 100644 index 0000000000..114d291968 --- /dev/null +++ b/source/tests/f.out @@ -0,0 +1 @@ +5.355088169393570574e+00,5.606772412401632266e+00,2.703270748296462966e-01,5.381408138049708967e+00,5.261355614357515087e+00,-4.079549918988088497e-01,-5.182324474551911919e+00,3.695481388907447262e-01,-5.238474288082559799e-02,1.665564584447352670e-01,-5.955401876564963892e+00,-2.217626865156164251e-01,-5.967343479332643419e+00,9.073821102416884665e-02,3.703103995504785639e-01,2.466151879965445548e-01,-5.373012500109096479e+00,4.146494691512619957e-02 diff --git a/source/tests/system/set.000/aparam.npy b/source/tests/system/set.000/aparam.npy new file mode 100644 index 0000000000000000000000000000000000000000..b3544695fa75af4593fe8b9c0fb99c6dffd39758 GIT binary patch literal 224 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= nXCxM+0{I$-Itpew3Pzeb3bhL40j^myL13pnggyzSsmBKZExl5t literal 0 HcmV?d00001 diff --git a/source/tests/system/set.000/box.npy b/source/tests/system/set.000/box.npy new file mode 100644 index 0000000000000000000000000000000000000000..f70e95b6e772550e087d1aba097f26b3b380704e GIT binary patch literal 200 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= hXCxM+0{I$-ItrGWItsN4WCJb+Ffeg|(6~s#@&G_89MJ#( literal 0 HcmV?d00001 diff --git a/source/tests/system/set.000/coord.npy b/source/tests/system/set.000/coord.npy new file mode 100644 index 0000000000000000000000000000000000000000..27a06e944251c351ad110c8c69cc764c2fcd84db GIT binary patch literal 272 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I$-ItqpsnmP)#3giN=`6m;X{f%5|zg$iK+xCJ@_C<$GYa?gvv_GBq Date: Tue, 27 Apr 2021 23:06:29 +0800 Subject: [PATCH 2/2] add type embedding method --- deepmd/train/trainer.py | 4 - source/tests/test_model_se_a_type.py | 117 +++++++++++++++++++++++++++ source/tests/type.out | 6 ++ source/tests/water_se_a_type.json | 58 +++++++++++++ 4 files changed, 181 insertions(+), 4 deletions(-) create mode 100644 source/tests/test_model_se_a_type.py create mode 100644 source/tests/type.out create mode 100644 source/tests/water_se_a_type.json diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index 2eebec1e8c..58bdbfbf2b 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -476,8 +476,6 @@ def train (self, train_data, valid_data=None) : tb_valid_writer = None train_time = 0 - varvar2 = self.sess.graph.get_tensor_by_name('t_embed:0') - print(self.sess.run(varvar2)) while cur_batch < stop_batch : # first round validation: @@ -515,8 +513,6 @@ def train (self, train_data, valid_data=None) : log.info("batch %7d training time %.2f s, testing time %.2f s" % (cur_batch, train_time, test_time)) train_time = 0 - varvar2 = self.sess.graph.get_tensor_by_name('t_embed:0') - print(self.sess.run(varvar2)) if self.save_freq > 0 and cur_batch % self.save_freq == 0 and self.run_opt.is_chief : if self.saver is not None : self.saver.save (self.sess, os.getcwd() + "/" + self.save_ckpt) diff --git a/source/tests/test_model_se_a_type.py b/source/tests/test_model_se_a_type.py new file mode 100644 index 0000000000..11e0c45fdf --- /dev/null +++ b/source/tests/test_model_se_a_type.py @@ -0,0 +1,117 @@ + +import dpdata,os,sys,unittest +import numpy as np +from deepmd.env import tf +import pickle +from common import Data,gen_data, j_loader + +from deepmd.utils.data_system import DataSystem +from deepmd.descriptor import DescrptSeAEbd_type +from deepmd.fit import EnerFitting +from deepmd.model import EnerModel +from deepmd.common import j_must_have + +GLOBAL_ENER_FLOAT_PRECISION = tf.float64 +GLOBAL_TF_FLOAT_PRECISION = tf.float64 +GLOBAL_NP_FLOAT_PRECISION = np.float64 + +class TestModel(unittest.TestCase): + def setUp(self) : + gen_data() + + def test_model(self): + jfile = 'water_se_a_type.json' + jdata = j_loader(jfile) + + systems = j_must_have(jdata, 'systems') + set_pfx = j_must_have(jdata, 'set_prefix') + batch_size = j_must_have(jdata, 'batch_size') + test_size = j_must_have(jdata, 'numb_test') + batch_size = 1 + test_size = 1 + stop_batch = j_must_have(jdata, 'stop_batch') + rcut = j_must_have (jdata['model']['descriptor'], 'rcut') + + data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt = None) + + test_data = data.get_test () + numb_test = 1 + + jdata['model']['descriptor'].pop('type', None) + descrpt = DescrptSeAEbd_type(**jdata['model']['descriptor']) + jdata['model']['fitting_net']['descrpt'] = descrpt + fitting = EnerFitting(**jdata['model']['fitting_net']) + model = EnerModel(descrpt, fitting) + + # model._compute_dstats([test_data['coord']], [test_data['box']], [test_data['type']], [test_data['natoms_vec']], [test_data['default_mesh']]) + input_data = {'coord' : [test_data['coord']], + 'box': [test_data['box']], + 'type': [test_data['type']], + 'natoms_vec' : [test_data['natoms_vec']], + 'default_mesh' : [test_data['default_mesh']] + } + model._compute_input_stat(input_data) + model.descrpt.bias_atom_e = data.compute_energy_shift() + + t_prop_c = tf.placeholder(tf.float32, [5], name='t_prop_c') + t_energy = tf.placeholder(GLOBAL_ENER_FLOAT_PRECISION, [None], name='t_energy') + t_force = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_force') + t_virial = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_virial') + t_atom_ener = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_atom_ener') + t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='i_coord') + t_type = tf.placeholder(tf.int32, [None], name='i_type') + t_natoms = tf.placeholder(tf.int32, [model.ntypes+2], name='i_natoms') + t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, 9], name='i_box') + t_mesh = tf.placeholder(tf.int32, [None], name='i_mesh') + is_training = tf.placeholder(tf.bool) + t_fparam = None + + model_pred \ + = model.build (t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + t_fparam, + suffix = "se_a_type", + reuse = False) + energy = model_pred['energy'] + force = model_pred['force'] + virial = model_pred['virial'] + atom_ener = model_pred['atom_ener'] + + feed_dict_test = {t_prop_c: test_data['prop_c'], + t_energy: test_data['energy'] [:numb_test], + t_force: np.reshape(test_data['force'] [:numb_test, :], [-1]), + t_virial: np.reshape(test_data['virial'] [:numb_test, :], [-1]), + t_atom_ener: np.reshape(test_data['atom_ener'][:numb_test, :], [-1]), + t_coord: np.reshape(test_data['coord'] [:numb_test, :], [-1]), + t_box: test_data['box'] [:numb_test, :], + t_type: np.reshape(test_data['type'] [:numb_test, :], [-1]), + t_natoms: test_data['natoms_vec'], + t_mesh: test_data['default_mesh'], + is_training: False} + + sess = tf.Session() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], + feed_dict = feed_dict_test) + + e = e.reshape([-1]) + f = f.reshape([-1]) + v = v.reshape([-1]) + + refe = [60.49433111978599] + reff = [ 0.10244746560517419 ,0.11277079482689496 , 0.003946400478907016 , 0.14144699199575358 , 0.13176375591292006, -0.012339293048477881,-0.11117992687685702 , 0.06595193763888685 , 0.0012527024605403633 , 0.07939089213976509 ,-0.18379039068094222 ,-0.006193350470194934,-0.27142564024593846 , 0.04836512503682785 , 0.014982194455435308 , 0.05932021738210261 ,-0.17506122273458752, -0.0016486538762098774] + refv = [-0.6950608872626888 , 0.14564405160365884 , 0.02145338180571397 , 0.14564405160365884 ,-0.3895317073251466, -0.00780801250611546,0.021453381805713958 ,-0.0078080125061154625 ,-0.0012301327522364124] + refe = np.reshape(refe, [-1]) + reff = np.reshape(reff, [-1]) + refv = np.reshape(refv, [-1]) + + places = 10 + for ii in range(e.size) : + self.assertAlmostEqual(e[ii], refe[ii], places = places) + for ii in range(f.size) : + self.assertAlmostEqual(f[ii], reff[ii], places = places) + for ii in range(v.size) : + self.assertAlmostEqual(v[ii], refv[ii], places = places) diff --git a/source/tests/type.out b/source/tests/type.out new file mode 100644 index 0000000000..4d33159df5 --- /dev/null +++ b/source/tests/type.out @@ -0,0 +1,6 @@ +[60.49433112] +[ 0.10244747 0.11277079 0.0039464 0.14144699 0.13176376 -0.01233929 + -0.11117993 0.06595194 0.0012527 0.07939089 -0.18379039 -0.00619335 + -0.27142564 0.04836513 0.01498219 0.05932022 -0.17506122 -0.00164865] +[-0.69506089 0.14564405 0.02145338 0.14564405 -0.38953171 -0.00780801 + 0.02145338 -0.00780801 -0.00123013] diff --git a/source/tests/water_se_a_type.json b/source/tests/water_se_a_type.json new file mode 100644 index 0000000000..792ffcf3b8 --- /dev/null +++ b/source/tests/water_se_a_type.json @@ -0,0 +1,58 @@ +{ + "_comment": " model parameters", + "model" : { + "descriptor" :{ + "type": "se_a_type", + "sel": [46, 92], + "rcut_smth": 5.80, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": false, + "type_filter": [5,10,10], + "type_one_side": false, + "axis_neuron": 16, + "seed": 1 + }, + "fitting_net" : { + "neuron": [240, 240, 240], + "resnet_dt": true, + "share_fitting": true, + "seed": 1 + } + }, + + + "_comment": " traing controls", + "systems": ["system"], + "set_prefix": "set", + "stop_batch": 1000000, + "batch_size": 1, + "start_lr": 0.005, + "decay_steps": 5000, + "decay_rate": 0.95, + + "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, + + "seed": 1, + + "_comment": " display and restart", + "_comment": " frequencies counted in batch", + "disp_file": "lcurve.out", + "disp_freq": 100, + "numb_test": 1, + "save_freq": 1000, + "save_ckpt": "model.ckpt", + "load_ckpt": "model.ckpt", + "disp_training": true, + "time_training": true, + "profiling": false, + "profiling_file": "timeline.json", + + "_comment": "that's all" +} +