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)