def build(self,
coord_,
atype_,
natoms,
box_,
mesh,
davg=None,
dstd=None,
suffix='',
reuse=None):
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)
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.descrpt_se_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)
self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt])
self.dout, self.qmat = self._pass_filter(self.descrpt_reshape,
natoms,
suffix=suffix,
reuse=reuse,
trainable=self.trainable)
return self.dout
def one_layer(inputs,
outputs_size,
activation_fn=tf.nn.tanh,
precision=global_tf_float_precision,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
use_timestep=False,
trainable=True,
useBN=False):
with tf.variable_scope(name, reuse=reuse):
shape = inputs.get_shape().as_list()
w = tf.get_variable(
'matrix', [shape[1], outputs_size],
precision,
tf.random_normal_initializer(stddev=stddev /
np.sqrt(shape[1] + outputs_size),
seed=seed),
trainable=trainable)
b = tf.get_variable('bias', [outputs_size],
precision,
tf.random_normal_initializer(stddev=stddev,
mean=bavg,
seed=seed),
trainable=trainable)
hidden = tf.matmul(inputs, w) + b
if activation_fn != None and use_timestep:
idt = tf.get_variable('idt', [outputs_size],
precision,
tf.random_normal_initializer(stddev=0.001,
mean=0.1,
seed=seed),
trainable=trainable)
if activation_fn != None:
if useBN:
None
# hidden_bn = self._batch_norm(hidden, name=name+'_normalization', reuse=reuse)
# return activation_fn(hidden_bn)
else:
if use_timestep:
return tf.reshape(activation_fn(hidden),
[-1, outputs_size]) * idt
else:
return tf.reshape(activation_fn(hidden),
[-1, outputs_size])
else:
if useBN:
None
# return self._batch_norm(hidden, name=name+'_normalization', reuse=reuse)
else:
return hidden
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]))
self.nei_type = tf.get_variable(
't_nei_type', [self.nnei],
dtype=GLOBAL_TF_FLOAT_PRECISION,
trainable=False,
initializer=tf.constant_initializer(nei_type))
self.dout = DescrptSeA.build(self,
coord_,
atype_,
natoms,
box_,
mesh,
input_dict,
suffix=suffix,
reuse=reuse)
tf.summary.histogram('embedding_net_output', self.dout)
return self.dout
def _net(self, inputs, name, reuse=False):
with tf.variable_scope(name, reuse=reuse):
net_w = tf.get_variable('net_w', [self.ndescrpt],
global_tf_float_precision,
tf.constant_initializer(self.net_w_i))
dot_v = tf.matmul(tf.reshape(inputs, [-1, self.ndescrpt]),
tf.reshape(net_w, [self.ndescrpt, 1]))
return tf.reshape(dot_v, [-1])
def _net(self, inputs, name, reuse=False):
with tf.variable_scope(name, reuse=reuse):
net_w = tf.get_variable('net_w', [self.ndescrpt],
GLOBAL_TF_FLOAT_PRECISION,
tf.constant_initializer(self.net_w_i))
dot_v = tf.matmul(tf.reshape(inputs, [-1, self.ndescrpt]),
tf.reshape(net_w, [self.ndescrpt, 1]))
return tf.reshape(dot_v, [-1])
def comp_f_dw(self, dcoord, dbox, dtype, tnatoms, name, reuse=None):
energy, force, virial = self.comp_ef(dcoord, dbox, dtype, tnatoms,
name, reuse)
with tf.variable_scope(name, reuse=True):
net_w = tf.get_variable('net_w', [self.ndescrpt],
global_tf_float_precision,
tf.constant_initializer(self.net_w_i))
f_mag = tf.reduce_sum(tf.nn.tanh(force))
f_mag_dw = tf.gradients(f_mag, net_w)
assert (len(f_mag_dw) == 1), "length of dw is wrong"
return f_mag, f_mag_dw[0]
def comp_v_dw(self, dcoord, dbox, dtype, tnatoms, name, reuse=None):
energy, force, virial = self.comp_ef(dcoord, dbox, dtype, tnatoms,
name, reuse)
with tf.variable_scope(name, reuse=True):
net_w = tf.get_variable('net_w', [self.ndescrpt],
GLOBAL_TF_FLOAT_PRECISION,
tf.constant_initializer(self.net_w_i))
v_mag = tf.reduce_sum(virial)
v_mag_dw = tf.gradients(v_mag, net_w)
assert (len(v_mag_dw) == 1), "length of dw is wrong"
return v_mag, v_mag_dw[0]
def __init__(self, data):
# tabulated
Inter.__init__(self, data)
_make_tab(data.get_ntypes())
self.srtab = TabInter('tab.xvg')
self.smin_alpha = 0.3
self.sw_rmin = 1
self.sw_rmax = 3.45
tab_info, tab_data = self.srtab.get()
with tf.variable_scope('tab', reuse=tf.AUTO_REUSE):
self.tab_info = tf.get_variable(
't_tab_info',
tab_info.shape,
dtype=tf.float64,
trainable=False,
initializer=tf.constant_initializer(tab_info))
self.tab_data = tf.get_variable(
't_tab_data',
tab_data.shape,
dtype=tf.float64,
trainable=False,
initializer=tf.constant_initializer(tab_data))
def build(self,
coord_,
atype_,
natoms,
box,
mesh,
input_dict,
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_mt = tf.constant(self.model_type,
name='model_type',
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]])
dout \
= self.descrpt.build(coord_,
atype_,
natoms,
box,
mesh,
davg = self.davg,
dstd = self.dstd,
suffix = suffix,
reuse = reuse)
dout = tf.identity(dout, name='o_descriptor')
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,
input_dict,
natoms,
bias_atom_e=self.bias_atom_e,
reuse=reuse,
suffix=suffix)
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.tab_inter(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
return model_dict
def _filter_type_ext(self,
inputs,
natoms,
activation_fn=tf.nn.tanh,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
trainable=True):
# natom x (nei x 4)
outputs_size = [1] + self.filter_neuron
outputs_size_2 = self.n_axis_neuron
with tf.variable_scope(name, reuse=reuse):
start_index = 0
result_all = []
xyz_scatter_1_all = []
xyz_scatter_2_all = []
for type_i in range(self.ntypes):
# cut-out inputs
# with natom x (nei_type_i x 4)
inputs_i = tf.slice(inputs, [0, start_index * 4],
[-1, self.sel_a[type_i] * 4])
start_index += self.sel_a[type_i]
shape_i = inputs_i.get_shape().as_list()
# with (natom x nei_type_i) x 4
inputs_reshape = tf.reshape(inputs_i, [-1, 4])
xyz_scatter = tf.reshape(
tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
for ii in range(1, len(outputs_size)):
w = tf.get_variable(
'matrix_' + str(ii) + '_' + str(type_i),
[outputs_size[ii - 1], outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(
stddev=stddev /
np.sqrt(outputs_size[ii] + outputs_size[ii - 1]),
seed=seed),
trainable=trainable)
b = tf.get_variable('bias_' + str(ii) + '_' + str(type_i),
[1, outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(
stddev=stddev,
mean=bavg,
seed=seed),
trainable=trainable)
if self.filter_resnet_dt:
idt = tf.get_variable(
'idt_' + str(ii) + '_' + str(type_i),
[1, outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(stddev=0.001,
mean=1.0,
seed=seed),
trainable=trainable)
if outputs_size[ii] == outputs_size[ii - 1]:
if self.filter_resnet_dt:
xyz_scatter += activation_fn(
tf.matmul(xyz_scatter, w) + b) * idt
else:
xyz_scatter += activation_fn(
tf.matmul(xyz_scatter, w) + b)
elif outputs_size[ii] == outputs_size[ii - 1] * 2:
if self.filter_resnet_dt:
xyz_scatter = tf.concat(
[xyz_scatter, xyz_scatter], 1) + activation_fn(
tf.matmul(xyz_scatter, w) + b) * idt
else:
xyz_scatter = tf.concat(
[xyz_scatter, xyz_scatter], 1) + activation_fn(
tf.matmul(xyz_scatter, w) + b)
else:
xyz_scatter = activation_fn(
tf.matmul(xyz_scatter, w) + b)
# natom x nei_type_i x out_size
xyz_scatter = tf.reshape(
xyz_scatter, (-1, shape_i[1] // 4, outputs_size[-1]))
# natom x nei_type_i x 4
inputs_i_reshape = tf.reshape(inputs_i,
[-1, shape_i[1] // 4, 4])
# natom x 4 x outputs_size
xyz_scatter_1 = tf.matmul(inputs_i_reshape,
xyz_scatter,
transpose_a=True)
xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape_i[1])
# 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_1_all.append(xyz_scatter_1)
xyz_scatter_2_all.append(xyz_scatter_2)
# for type_i in range(self.ntypes):
# for type_j in range(type_i, self.ntypes):
# # natom x outputs_size x outputs_size_2
# result = tf.matmul(xyz_scatter_1_all[type_i], xyz_scatter_2_all[type_j], transpose_a = True)
# # natom x (outputs_size x outputs_size_2)
# result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]])
# result_all.append(tf.identity(result))
xyz_scatter_2_coll = tf.concat(xyz_scatter_2_all, axis=2)
for type_i in range(self.ntypes):
# natom x outputs_size x (outputs_size_2 x ntypes)
result = tf.matmul(xyz_scatter_1_all[type_i],
xyz_scatter_2_coll,
transpose_a=True)
# natom x (outputs_size x outputs_size_2 x ntypes)
result = tf.reshape(
result,
[-1, outputs_size_2 * self.ntypes * outputs_size[-1]])
result_all.append(tf.identity(result))
# natom x (ntypes x outputs_size x outputs_size_2 x ntypes)
result_all = tf.concat(result_all, axis=1)
return result_all
def _filter(self,
inputs,
type_input,
natoms,
activation_fn=tf.nn.tanh,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
trainable=True):
# natom x (nei x 4)
shape = inputs.get_shape().as_list()
outputs_size = [1] + self.filter_neuron
outputs_size_2 = self.n_axis_neuron
with tf.variable_scope(name, reuse=reuse):
start_index = 0
xyz_scatter_total = []
for type_i in range(self.ntypes):
# cut-out inputs
# with natom x (nei_type_i x 4)
inputs_i = tf.slice(inputs, [0, start_index * 4],
[-1, self.sel_a[type_i] * 4])
start_index += self.sel_a[type_i]
shape_i = inputs_i.get_shape().as_list()
# with (natom x nei_type_i) x 4
inputs_reshape = tf.reshape(inputs_i, [-1, 4])
xyz_scatter = tf.reshape(
tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
if (type_input, type_i) not in self.exclude_types:
for ii in range(1, len(outputs_size)):
w = tf.get_variable(
'matrix_' + str(ii) + '_' + str(type_i),
[outputs_size[ii - 1], outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(
stddev=stddev / np.sqrt(outputs_size[ii] +
outputs_size[ii - 1]),
seed=seed),
trainable=trainable)
b = tf.get_variable(
'bias_' + str(ii) + '_' + str(type_i),
[1, outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(stddev=stddev,
mean=bavg,
seed=seed),
trainable=trainable)
if self.filter_resnet_dt:
idt = tf.get_variable(
'idt_' + str(ii) + '_' + str(type_i),
[1, outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(stddev=0.001,
mean=1.0,
seed=seed),
trainable=trainable)
if outputs_size[ii] == outputs_size[ii - 1]:
if self.filter_resnet_dt:
xyz_scatter += activation_fn(
tf.matmul(xyz_scatter, w) + b) * idt
else:
xyz_scatter += activation_fn(
tf.matmul(xyz_scatter, w) + b)
elif outputs_size[ii] == outputs_size[ii - 1] * 2:
if self.filter_resnet_dt:
xyz_scatter = tf.concat(
[xyz_scatter, xyz_scatter],
1) + activation_fn(
tf.matmul(xyz_scatter, w) + b) * idt
else:
xyz_scatter = tf.concat(
[xyz_scatter, xyz_scatter],
1) + activation_fn(
tf.matmul(xyz_scatter, w) + b)
else:
xyz_scatter = activation_fn(
tf.matmul(xyz_scatter, w) + b)
else:
w = tf.zeros((outputs_size[0], outputs_size[-1]),
dtype=global_tf_float_precision)
xyz_scatter = tf.matmul(xyz_scatter, w)
# natom x nei_type_i 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)
# 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,
xyz_scatter,
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[-1]])
return result, qmat
def one_layer(inputs,
outputs_size,
activation_fn=tf.nn.tanh,
precision = GLOBAL_TF_FLOAT_PRECISION,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
use_timestep = False,
trainable = True,
useBN = False,
uniform_seed = False,
initial_variables = None):
with tf.variable_scope(name, reuse=reuse):
shape = inputs.get_shape().as_list()
w_initializer = tf.random_normal_initializer(
stddev=stddev / np.sqrt(shape[1] + outputs_size),
seed=seed if (seed is None or uniform_seed) else seed + 0)
b_initializer = tf.random_normal_initializer(
stddev=stddev,
mean=bavg,
seed=seed if (seed is None or uniform_seed) else seed + 1)
if initial_variables is not None:
w_initializer = tf.constant_initializer(initial_variables[name + '/matrix'])
b_initializer = tf.constant_initializer(initial_variables[name + '/bias'])
w = tf.get_variable('matrix',
[shape[1], outputs_size],
precision,
w_initializer,
trainable = trainable)
variable_summaries(w, 'matrix')
b = tf.get_variable('bias',
[outputs_size],
precision,
b_initializer,
trainable = trainable)
variable_summaries(b, 'bias')
hidden = tf.matmul(inputs, w) + b
if activation_fn != None and use_timestep :
idt_initializer = tf.random_normal_initializer(
stddev=0.001,
mean=0.1,
seed=seed if (seed is None or uniform_seed) else seed + 2)
if initial_variables is not None:
idt_initializer = tf.constant_initializer(initial_variables[name + '/idt'])
idt = tf.get_variable('idt',
[outputs_size],
precision,
idt_initializer,
trainable = trainable)
variable_summaries(idt, 'idt')
if activation_fn != None:
if useBN:
None
# hidden_bn = self._batch_norm(hidden, name=name+'_normalization', reuse=reuse)
# return activation_fn(hidden_bn)
else:
if use_timestep :
return tf.reshape(activation_fn(hidden), [-1, outputs_size]) * idt
else :
return tf.reshape(activation_fn(hidden), [-1, outputs_size])
else:
if useBN:
None
# return self._batch_norm(hidden, name=name+'_normalization', reuse=reuse)
else:
return hidden
def _filter_r(self,
inputs,
type_input,
natoms,
activation_fn=tf.nn.tanh,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
trainable=True):
# natom x nei
outputs_size = [1] + self.filter_neuron
with tf.variable_scope(name, reuse=reuse):
start_index = 0
xyz_scatter_total = []
for type_i in range(self.ntypes):
# cut-out inputs
# with natom x nei_type_i
inputs_i = tf.slice(inputs, [0, start_index],
[-1, self.sel_r[type_i]])
start_index += self.sel_r[type_i]
shape_i = inputs_i.get_shape().as_list()
# with (natom x nei_type_i) x 1
xyz_scatter = tf.reshape(inputs_i, [-1, 1])
if (type_input, type_i) not in self.exclude_types:
for ii in range(1, len(outputs_size)):
w = tf.get_variable(
'matrix_' + str(ii) + '_' + str(type_i),
[outputs_size[ii - 1], outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(
stddev=stddev / np.sqrt(outputs_size[ii] +
outputs_size[ii - 1]),
seed=seed),
trainable=trainable)
b = tf.get_variable(
'bias_' + str(ii) + '_' + str(type_i),
[1, outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(stddev=stddev,
mean=bavg,
seed=seed),
trainable=trainable)
if self.filter_resnet_dt:
idt = tf.get_variable(
'idt_' + str(ii) + '_' + str(type_i),
[1, outputs_size[ii]],
self.filter_precision,
tf.random_normal_initializer(stddev=0.001,
mean=1.0,
seed=seed),
trainable=trainable)
if outputs_size[ii] == outputs_size[ii - 1]:
if self.filter_resnet_dt:
xyz_scatter += activation_fn(
tf.matmul(xyz_scatter, w) + b) * idt
else:
xyz_scatter += activation_fn(
tf.matmul(xyz_scatter, w) + b)
elif outputs_size[ii] == outputs_size[ii - 1] * 2:
if self.filter_resnet_dt:
xyz_scatter = tf.concat(
[xyz_scatter, xyz_scatter],
1) + activation_fn(
tf.matmul(xyz_scatter, w) + b) * idt
else:
xyz_scatter = tf.concat(
[xyz_scatter, xyz_scatter],
1) + activation_fn(
tf.matmul(xyz_scatter, w) + b)
else:
xyz_scatter = activation_fn(
tf.matmul(xyz_scatter, w) + b)
else:
w = tf.zeros((outputs_size[0], outputs_size[-1]),
dtype=global_tf_float_precision)
xyz_scatter = tf.matmul(xyz_scatter, w)
# natom x nei_type_i x out_size
xyz_scatter = tf.reshape(xyz_scatter,
(-1, shape_i[1], outputs_size[-1]))
xyz_scatter_total.append(xyz_scatter)
# natom x nei x outputs_size
xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
# natom x outputs_size
#
res_rescale = 1. / 5.
result = tf.reduce_mean(xyz_scatter, axis=1) * res_rescale
return result
def build(
self,
inputs: tf.Tensor,
natoms: tf.Tensor,
input_dict: dict = None,
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
Returns
-------
ener
The system energy
"""
if input_dict is None:
input_dict = {}
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))
inputs = tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]])
if len(self.atom_ener):
# only for atom_ener
nframes = input_dict.get('nframes')
if nframes is not None:
# like inputs, but we don't want to add a dependency on inputs
inputs_zero = tf.zeros((nframes, self.dim_descrpt * natoms[0]),
dtype=self.fitting_precision)
else:
inputs_zero = tf.zeros_like(inputs,
dtype=self.fitting_precision)
if bias_atom_e is not None:
assert (len(bias_atom_e) == self.ntypes)
fparam = None
aparam = None
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]])
type_embedding = input_dict.get('type_embedding', None)
if type_embedding is not None:
atype_embed = embed_atom_type(self.ntypes, natoms, type_embedding)
atype_embed = tf.tile(atype_embed, [tf.shape(inputs)[0], 1])
else:
atype_embed = None
if atype_embed is None:
start_index = 0
outs_list = []
for type_i in range(self.ntypes):
if bias_atom_e is None:
type_bias_ae = 0.0
else:
type_bias_ae = bias_atom_e[type_i]
final_layer = self._build_lower(start_index,
natoms[2 + type_i],
inputs,
fparam,
aparam,
bias_atom_e=type_bias_ae,
suffix='_type_' + str(type_i) +
suffix,
reuse=reuse)
# concat the results
if type_i < len(
self.atom_ener) and self.atom_ener[type_i] is not None:
zero_layer = self._build_lower(start_index,
natoms[2 + type_i],
inputs_zero,
fparam,
aparam,
bias_atom_e=type_bias_ae,
suffix='_type_' +
str(type_i) + suffix,
reuse=True)
final_layer += self.atom_ener[type_i] - zero_layer
final_layer = tf.reshape(
final_layer, [tf.shape(inputs)[0], natoms[2 + type_i]])
outs_list.append(final_layer)
start_index += natoms[2 + type_i]
# concat the results
# concat once may be faster than multiple concat
outs = tf.concat(outs_list, axis=1)
# with type embedding
else:
if len(self.atom_ener) > 0:
raise RuntimeError(
"setting atom_ener is not supported by type embedding")
atype_embed = tf.cast(atype_embed, self.fitting_precision)
type_shape = atype_embed.get_shape().as_list()
inputs = tf.concat(
[tf.reshape(inputs, [-1, self.dim_descrpt]), atype_embed],
axis=1)
self.dim_descrpt = self.dim_descrpt + type_shape[1]
inputs = tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]])
final_layer = self._build_lower(0,
natoms[0],
inputs,
fparam,
aparam,
bias_atom_e=0.0,
suffix=suffix,
reuse=reuse)
outs = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[0]])
# add atom energy bias; TF will broadcast to all batches
# tf.repeat is avaiable in TF>=2.1 or TF 1.15
_TF_VERSION = Version(TF_VERSION)
if (Version('1.15') <= _TF_VERSION < Version('2') or _TF_VERSION >=
Version('2.1')) and self.bias_atom_e is not None:
outs += tf.repeat(
tf.Variable(self.bias_atom_e,
dtype=self.fitting_precision,
trainable=False,
name="bias_atom_ei"), natoms[2:])
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.reshape(outs, [-1])
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
"""
davg = self.davg
dstd = self.dstd
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(self.rcut,
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_r(coord,
atype,
natoms,
box,
mesh,
self.t_avg,
self.t_std,
rcut = self.rcut,
rcut_smth = self.rcut_smth,
sel = self.sel_r)
self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt])
self._identity_tensors(suffix=suffix)
# 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.dout = self._pass_filter(self.descrpt_reshape,
natoms,
suffix=suffix,
reuse=reuse,
trainable=self.trainable)
tf.summary.histogram('embedding_net_output', self.dout)
return self.dout
def build(self,
coord_,
atype_,
natoms,
box_,
mesh,
input_dict,
suffix='',
reuse=None):
efield = input_dict['efield']
davg = self.davg
dstd = self.dstd
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]])
efield = tf.reshape(efield, [-1, 3])
efield = self._normalize_3d(efield)
efield = tf.reshape(efield, [-1, natoms[0] * 3])
self.descrpt, self.descrpt_deriv, self.rij, self.nlist \
= self.op (coord,
atype,
natoms,
box,
mesh,
efield,
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)
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')
# 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.dout, self.qmat = self._pass_filter(self.descrpt_reshape,
atype,
natoms,
input_dict,
suffix=suffix,
reuse=reuse,
trainable=self.trainable)
tf.summary.histogram('embedding_net_output', self.dout)
return self.dout
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_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]])
# 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
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)
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 one_layer(inputs,
outputs_size,
activation_fn=tf.nn.tanh,
precision=GLOBAL_TF_FLOAT_PRECISION,
stddev=1.0,
bavg=0.0,
name='linear',
reuse=None,
seed=None,
use_timestep=False,
trainable=True,
useBN=False,
uniform_seed=False,
initial_variables=None,
mixed_prec=None,
final_layer=False):
# For good accuracy, the last layer of the fitting network uses a higher precision neuron network.
if mixed_prec is not None and final_layer:
inputs = tf.cast(inputs, get_precision(mixed_prec['output_prec']))
with tf.variable_scope(name, reuse=reuse):
shape = inputs.get_shape().as_list()
w_initializer = tf.random_normal_initializer(
stddev=stddev / np.sqrt(shape[1] + outputs_size),
seed=seed if (seed is None or uniform_seed) else seed + 0)
b_initializer = tf.random_normal_initializer(
stddev=stddev,
mean=bavg,
seed=seed if (seed is None or uniform_seed) else seed + 1)
if initial_variables is not None:
w_initializer = tf.constant_initializer(
initial_variables[name + '/matrix'])
b_initializer = tf.constant_initializer(initial_variables[name +
'/bias'])
w = tf.get_variable('matrix', [shape[1], outputs_size],
precision,
w_initializer,
trainable=trainable)
variable_summaries(w, 'matrix')
b = tf.get_variable('bias', [outputs_size],
precision,
b_initializer,
trainable=trainable)
variable_summaries(b, 'bias')
if mixed_prec is not None and not final_layer:
inputs = tf.cast(inputs, get_precision(mixed_prec['compute_prec']))
w = tf.cast(w, get_precision(mixed_prec['compute_prec']))
b = tf.cast(b, get_precision(mixed_prec['compute_prec']))
hidden = tf.nn.bias_add(tf.matmul(inputs, w), b)
if activation_fn != None and use_timestep:
idt_initializer = tf.random_normal_initializer(
stddev=0.001,
mean=0.1,
seed=seed if (seed is None or uniform_seed) else seed + 2)
if initial_variables is not None:
idt_initializer = tf.constant_initializer(
initial_variables[name + '/idt'])
idt = tf.get_variable('idt', [outputs_size],
precision,
idt_initializer,
trainable=trainable)
variable_summaries(idt, 'idt')
if activation_fn != None:
if useBN:
None
# hidden_bn = self._batch_norm(hidden, name=name+'_normalization', reuse=reuse)
# return activation_fn(hidden_bn)
else:
if use_timestep:
if mixed_prec is not None and not final_layer:
idt = tf.cast(
idt, get_precision(mixed_prec['compute_prec']))
hidden = tf.reshape(activation_fn(hidden),
[-1, outputs_size]) * idt
else:
hidden = tf.reshape(activation_fn(hidden),
[-1, outputs_size])
if mixed_prec is not None:
hidden = tf.cast(hidden, get_precision(mixed_prec['output_prec']))
return hidden
def embedding_net(xx,
network_size,
precision,
activation_fn=tf.nn.tanh,
resnet_dt=False,
name_suffix='',
stddev=1.0,
bavg=0.0,
seed=None,
trainable=True,
uniform_seed=False,
initial_variables=None,
mixed_prec=None):
r"""The embedding network.
The embedding network function :math:`\mathcal{N}` is constructed by is the
composition of multiple layers :math:`\mathcal{L}^{(i)}`:
.. math::
\mathcal{N} = \mathcal{L}^{(n)} \circ \mathcal{L}^{(n-1)}
\circ \cdots \circ \mathcal{L}^{(1)}
A layer :math:`\mathcal{L}` is given by one of the following forms,
depending on the number of nodes: [1]_
.. math::
\mathbf{y}=\mathcal{L}(\mathbf{x};\mathbf{w},\mathbf{b})=
\begin{cases}
\boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b}) + \mathbf{x}, & N_2=N_1 \\
\boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b}) + (\mathbf{x}, \mathbf{x}), & N_2 = 2N_1\\
\boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b}), & \text{otherwise} \\
\end{cases}
where :math:`\mathbf{x} \in \mathbb{R}^{N_1}`$` is the input vector and :math:`\mathbf{y} \in \mathbb{R}^{N_2}`
is the output vector. :math:`\mathbf{w} \in \mathbb{R}^{N_1 \times N_2}` and
:math:`\mathbf{b} \in \mathbb{R}^{N_2}`$` are weights and biases, respectively,
both of which are trainable if `trainable` is `True`. :math:`\boldsymbol{\phi}`
is the activation function.
Parameters
----------
xx : Tensor
Input tensor :math:`\mathbf{x}` 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 :math:`\boldsymbol{\phi}`
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 network is trainable
uniform_seed : boolean
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
initial_variables : dict
The input dict which stores the embedding net variables
mixed_prec
The input dict which stores the mixed precision setting for the embedding net
References
----------
.. [1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Identitymappings
in deep residual networks. InComputer Vision – ECCV 2016,pages 630–645. Springer
International Publishing, 2016.
"""
input_shape = xx.get_shape().as_list()
outputs_size = [input_shape[1]] + network_size
for ii in range(1, len(outputs_size)):
w_initializer = tf.random_normal_initializer(
stddev=stddev / np.sqrt(outputs_size[ii] + outputs_size[ii - 1]),
seed=seed if (seed is None or uniform_seed) else seed + ii * 3 + 0)
b_initializer = tf.random_normal_initializer(
stddev=stddev,
mean=bavg,
seed=seed if (seed is None or uniform_seed) else seed + 3 * ii + 1)
if initial_variables is not None:
scope = tf.get_variable_scope().name
w_initializer = tf.constant_initializer(
initial_variables[scope + '/matrix_' + str(ii) + name_suffix])
b_initializer = tf.constant_initializer(
initial_variables[scope + '/bias_' + str(ii) + name_suffix])
w = tf.get_variable('matrix_' + str(ii) + name_suffix,
[outputs_size[ii - 1], outputs_size[ii]],
precision,
w_initializer,
trainable=trainable)
variable_summaries(w, 'matrix_' + str(ii) + name_suffix)
b = tf.get_variable('bias_' + str(ii) + name_suffix,
[outputs_size[ii]],
precision,
b_initializer,
trainable=trainable)
variable_summaries(b, 'bias_' + str(ii) + name_suffix)
if mixed_prec is not None:
xx = tf.cast(xx, get_precision(mixed_prec['compute_prec']))
w = tf.cast(w, get_precision(mixed_prec['compute_prec']))
b = tf.cast(b, get_precision(mixed_prec['compute_prec']))
hidden = tf.reshape(activation_fn(tf.nn.bias_add(tf.matmul(xx, w), b)),
[-1, outputs_size[ii]])
if resnet_dt:
idt_initializer = tf.random_normal_initializer(
stddev=0.001,
mean=1.0,
seed=seed if
(seed is None or uniform_seed) else seed + 3 * ii + 2)
if initial_variables is not None:
scope = tf.get_variable_scope().name
idt_initializer = tf.constant_initializer(
initial_variables[scope + '/idt_' + str(ii) + name_suffix])
idt = tf.get_variable('idt_' + str(ii) + name_suffix,
[1, outputs_size[ii]],
precision,
idt_initializer,
trainable=trainable)
variable_summaries(idt, 'idt_' + str(ii) + name_suffix)
if mixed_prec is not None:
idt = tf.cast(idt, get_precision(mixed_prec['compute_prec']))
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
if mixed_prec is not None:
xx = tf.cast(xx, get_precision(mixed_prec['output_prec']))
return xx
def build(self, inputs, input_dict, natoms, reuse=None, suffix=''):
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))
start_index = 0
inputs = tf.cast(
tf.reshape(inputs, [-1, self.dim_descrpt * 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]])
for type_i in range(self.ntypes):
# cut-out inputs
inputs_i = tf.slice(inputs, [0, start_index * self.dim_descrpt],
[-1, natoms[2 + type_i] * self.dim_descrpt])
inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
layer = inputs_i
if self.numb_fparam > 0:
ext_fparam = tf.tile(fparam, [1, natoms[2 + type_i]])
ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam])
layer = tf.concat([layer, ext_fparam], axis=1)
if self.numb_aparam > 0:
ext_aparam = tf.slice(
aparam, [0, start_index * self.numb_aparam],
[-1, natoms[2 + type_i] * self.numb_aparam])
ext_aparam = tf.reshape(ext_aparam, [-1, self.numb_aparam])
layer = tf.concat([layer, ext_aparam], axis=1)
start_index += natoms[2 + type_i]
if bias_atom_e is None:
type_bias_ae = 0.0
else:
type_bias_ae = bias_atom_e[type_i]
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='layer_' + str(ii) + '_type_' + str(type_i) +
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='layer_' + str(ii) + '_type_' +
str(type_i) + 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='final_layer_type_' + str(type_i) +
suffix,
reuse=reuse,
seed=self.seed,
precision=self.fitting_precision,
trainable=self.trainable[-1])
if type_i < len(
self.atom_ener) and self.atom_ener[type_i] is not None:
inputs_zero = tf.zeros_like(inputs_i,
dtype=global_tf_float_precision)
layer = inputs_zero
if self.numb_fparam > 0:
layer = tf.concat([layer, ext_fparam], axis=1)
if self.numb_aparam > 0:
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='layer_' + str(ii) + '_type_' + str(type_i) +
suffix,
reuse=True,
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='layer_' + str(ii) + '_type_' + str(type_i) +
suffix,
reuse=True,
seed=self.seed,
activation_fn=self.fitting_activation_fn,
precision=self.fitting_precision,
trainable=self.trainable[ii])
zero_layer = one_layer(layer,
1,
activation_fn=None,
bavg=type_bias_ae,
name='final_layer_type_' + str(type_i) +
suffix,
reuse=True,
seed=self.seed,
precision=self.fitting_precision,
trainable=self.trainable[-1])
final_layer += self.atom_ener[type_i] - zero_layer
final_layer = tf.reshape(final_layer,
[tf.shape(inputs)[0], natoms[2 + type_i]])
# concat the results
if type_i == 0:
outs = final_layer
else:
outs = tf.concat([outs, final_layer], axis=1)
return tf.cast(tf.reshape(outs, [-1]), global_tf_float_precision)
