def _do_work(jdata, run_opt):
# init the model
model = NNPTrainer (jdata, run_opt = run_opt)
rcut = model.model.get_rcut()
type_map = model.model.get_type_map()
# init params and run options
assert('training' in jdata)
systems = j_must_have(jdata['training'], 'systems')
set_pfx = j_must_have(jdata['training'], 'set_prefix')
numb_sys = len(systems)
seed = None
if 'seed' in jdata['training'].keys() : seed = jdata['training']['seed']
if seed is not None:
seed = seed % (2**32)
np.random.seed (seed)
batch_size = j_must_have(jdata['training'], 'batch_size')
test_size = j_must_have(jdata['training'], 'numb_test')
stop_batch = j_must_have(jdata['training'], 'stop_batch')
if len(type_map) == 0:
# empty type_map
ipt_type_map = None
else:
ipt_type_map = type_map
data = DeepmdDataSystem(systems, batch_size, test_size, rcut, set_prefix=set_pfx, run_opt=run_opt, type_map = ipt_type_map)
data.add_dict(data_requirement)
# build the model with stats from the first system
model.build (data)
# train the model with the provided systems in a cyclic way
start_time = time.time()
cur_batch = 0
model.train (data, stop_batch)
end_time = time.time()
run_opt.message("finished training\nwall time: %.3f s" % (end_time-start_time))
def _do_work(jdata, run_opt):
# init the model
model = NNPTrainer(jdata, run_opt=run_opt)
rcut = model.model.get_rcut()
type_map = model.model.get_type_map()
# init params and run options
assert ('training' in jdata)
systems = j_must_have(jdata['training'], 'systems')
if type(systems) == str:
systems = expand_sys_str(systems)
set_pfx = j_must_have(jdata['training'], 'set_prefix')
seed = None
if 'seed' in jdata['training'].keys(): seed = jdata['training']['seed']
if seed is not None:
seed = seed % (2**32)
np.random.seed(seed)
batch_size = j_must_have(jdata['training'], 'batch_size')
test_size = j_must_have(jdata['training'], 'numb_test')
stop_batch = j_must_have(jdata['training'], 'stop_batch')
sys_probs = jdata['training'].get('sys_probs')
auto_prob_style = jdata['training'].get('auto_prob_style', 'prob_sys_size')
if len(type_map) == 0:
# empty type_map
ipt_type_map = None
else:
ipt_type_map = type_map
# data modifier
modifier = None
modi_data = jdata['model'].get("modifier", None)
if modi_data is not None:
if modi_data['type'] == 'dipole_charge':
modifier = DipoleChargeModifier(modi_data['model_name'],
modi_data['model_charge_map'],
modi_data['sys_charge_map'],
modi_data['ewald_h'],
modi_data['ewald_beta'])
else:
raise RuntimeError('unknown modifier type ' +
str(modi_data['type']))
# init data
data = DeepmdDataSystem(systems,
batch_size,
test_size,
rcut,
set_prefix=set_pfx,
type_map=ipt_type_map,
modifier=modifier)
data.print_summary(run_opt,
sys_probs=sys_probs,
auto_prob_style=auto_prob_style)
data.add_dict(data_requirement)
# build the model with stats from the first system
model.build(data, stop_batch)
# train the model with the provided systems in a cyclic way
start_time = time.time()
model.train(data)
end_time = time.time()
run_opt.message("finished training\nwall time: %.3f s" %
(end_time - start_time))
class Benchmark:
def __init__(self):
# setup
with open(os.path.join(os.path.dirname(__file__),
'water_se_a.json')) as fp:
jdata = json.load(fp)
self.run_opt = RunOptions(None)
self.run_opt.verbose = False
self.model = NNPTrainer(jdata, run_opt=self.run_opt)
rcut = self.model.model.get_rcut()
type_map = self.model.model.get_type_map()
systems = [os.path.join(os.path.dirname(__file__), 'data')]
set_pfx = jdata['training']['set_prefix']
seed = jdata['training']['seed']
np.random.seed(seed)
batch_size = jdata['training']['batch_size']
test_size = jdata['training']['numb_test']
self.data = DeepmdDataSystem(systems,
batch_size,
test_size,
rcut,
set_prefix=set_pfx,
run_opt=self.run_opt,
type_map=type_map)
self.data.add_dict(data_requirement)
self.model.build(self.data)
self.model._init_sess_serial()
cur_batch = self.model.sess.run(self.model.global_step)
self.cur_batch = cur_batch
def train_step(self):
batch_data = self.data.get_batch(sys_weights=self.model.sys_weights)
feed_dict_batch = {}
for kk in batch_data.keys():
if kk == 'find_type' or kk == 'type':
continue
if 'find_' in kk:
feed_dict_batch[self.model.place_holders[kk]] = batch_data[kk]
else:
feed_dict_batch[self.model.place_holders[kk]] = np.reshape(
batch_data[kk], [-1])
for ii in ['type']:
feed_dict_batch[self.model.place_holders[ii]] = np.reshape(
batch_data[ii], [-1])
for ii in ['natoms_vec', 'default_mesh']:
feed_dict_batch[self.model.place_holders[ii]] = batch_data[ii]
feed_dict_batch[self.model.place_holders['is_training']] = True
self.model.sess.run([self.model.train_op],
feed_dict=feed_dict_batch,
options=None,
run_metadata=None)
self.model.sess.run(self.model.global_step)
def _setUp(self):
args = Args()
run_opt = RunOptions(args, False)
with open (args.INPUT, 'r') as fp:
jdata = json.load (fp)
# init model
model = NNPTrainer (jdata, run_opt = run_opt)
rcut = model.model.get_rcut()
# init data system
systems = j_must_have(jdata['training'], 'systems')
set_pfx = j_must_have(jdata['training'], 'set_prefix')
batch_size = j_must_have(jdata['training'], 'batch_size')
test_size = j_must_have(jdata['training'], 'numb_test')
data = DeepmdDataSystem(systems,
batch_size,
test_size,
rcut,
set_prefix=set_pfx)
data.add_dict(data_requirement)
# clear the default graph
tf.reset_default_graph()
# build the model with stats from the first system
model.build (data)
# freeze the graph
with tf.Session() as sess:
init_op = tf.global_variables_initializer()
sess.run(init_op)
graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
nodes = "o_dipole,o_rmat,o_rmat_deriv,o_nlist,o_rij,descrpt_attr/rcut,descrpt_attr/ntypes,descrpt_attr/sel,descrpt_attr/ndescrpt,model_attr/tmap,model_attr/sel_type,model_attr/model_type"
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess,
input_graph_def,
nodes.split(",")
)
output_graph = os.path.join(modifier_datapath, 'dipole.pb')
with tf.gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
