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 test_get_batch(self):
batch_size = 3
test_size = 2
ds = DeepmdDataSystem(self.sys_name, batch_size, test_size, 2.0)
ds.add('test', self.test_ndof, atomic=True, must=True)
ds.add('null', self.test_ndof, atomic=True, must=False)
sys_idx = 0
data = ds.get_batch(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]),
list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_0/set.000/coord.npy'),
ds.get_sys(sys_idx).idx_map, 3, data['coord'])
self._in_array(np.load('sys_0/set.000/test.npy'),
ds.get_sys(sys_idx).idx_map, self.test_ndof,
data['test'])
self.assertAlmostEqual(
np.linalg.norm(
np.zeros([batch_size, self.natoms[sys_idx] * self.test_ndof]) -
data['null']), 0.0)
data = ds.get_batch(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]),
list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_0/set.001/coord.npy'),
ds.get_sys(sys_idx).idx_map, 3, data['coord'])
self._in_array(np.load('sys_0/set.001/test.npy'),
ds.get_sys(sys_idx).idx_map, self.test_ndof,
data['test'])
self.assertAlmostEqual(
np.linalg.norm(
np.zeros([batch_size, self.natoms[sys_idx] * self.test_ndof]) -
data['null']), 0.0)
data = ds.get_batch(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]),
list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_0/set.000/coord.npy'),
ds.get_sys(sys_idx).idx_map, 3, data['coord'])
self._in_array(np.load('sys_0/set.000/test.npy'),
ds.get_sys(sys_idx).idx_map, self.test_ndof,
data['test'])
self.assertAlmostEqual(
np.linalg.norm(
np.zeros([batch_size, self.natoms[sys_idx] * self.test_ndof]) -
data['null']), 0.0)
sys_idx = 2
data = ds.get_batch(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]),
list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_2/set.000/coord.npy'),
ds.get_sys(sys_idx).idx_map, 3, data['coord'])
self._in_array(np.load('sys_2/set.000/test.npy'),
ds.get_sys(sys_idx).idx_map, self.test_ndof,
data['test'])
self.assertAlmostEqual(
np.linalg.norm(
np.zeros([batch_size, self.natoms[sys_idx] * self.test_ndof]) -
data['null']), 0.0)
data = ds.get_batch(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]),
list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_2/set.001/coord.npy'),
ds.get_sys(sys_idx).idx_map, 3, data['coord'])
self._in_array(np.load('sys_2/set.001/test.npy'),
ds.get_sys(sys_idx).idx_map, self.test_ndof,
data['test'])
self.assertAlmostEqual(
np.linalg.norm(
np.zeros([batch_size, self.natoms[sys_idx] * self.test_ndof]) -
data['null']), 0.0)
data = ds.get_batch(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]),
list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_2/set.001/coord.npy'),
ds.get_sys(sys_idx).idx_map, 3, data['coord'])
self._in_array(np.load('sys_2/set.001/test.npy'),
ds.get_sys(sys_idx).idx_map, self.test_ndof,
data['test'])
self.assertAlmostEqual(
np.linalg.norm(
np.zeros([batch_size, self.natoms[sys_idx] * self.test_ndof]) -
data['null']), 0.0)
data = ds.get_batch(sys_idx=sys_idx)
self.assertEqual(list(data['type'][0]),
list(np.sort(self.atom_type[sys_idx])))
self._in_array(np.load('sys_2/set.000/coord.npy'),
ds.get_sys(sys_idx).idx_map, 3, data['coord'])
self._in_array(np.load('sys_2/set.000/test.npy'),
ds.get_sys(sys_idx).idx_map, self.test_ndof,
data['test'])
self.assertAlmostEqual(
np.linalg.norm(
np.zeros([batch_size, self.natoms[sys_idx] * self.test_ndof]) -
data['null']), 0.0)