diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py
index bf88a9a1e6..e5d291e7dc 100644
--- a/deepmd/entrypoints/main.py
+++ b/deepmd/entrypoints/main.py
@@ -232,10 +232,10 @@ def parse_args(args: Optional[List[str]] = None):
)
parser_tst.add_argument(
"-a",
- "--atomic-energy",
+ "--atomic",
action="store_true",
default=False,
- help="Test the accuracy of atomic energy",
+ help="Test the accuracy of atomic label, i.e. energy / dipole / polar",
)
# * compress model *****************************************************************
diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py
index 5c91547866..62687866a8 100644
--- a/deepmd/entrypoints/test.py
+++ b/deepmd/entrypoints/test.py
@@ -26,7 +26,7 @@ def test(
rand_seed: Optional[int],
shuffle_test: bool,
detail_file: str,
- atomic_energy: bool,
+ atomic: bool,
**kwargs,
):
"""Test model predictions.
@@ -47,7 +47,7 @@ def test(
whether to shuffle tests
detail_file : Optional[str]
file where test details will be output
- atomic_energy : bool
+ atomic : bool
whether per atom quantities should be computed
Raises
@@ -83,11 +83,11 @@ def test(
system,
numb_test,
detail_file,
- atomic_energy,
+ atomic,
append_detail=(cc != 0),
)
elif dp.model_type == "dipole":
- err, siz = test_dipole(dp, data, numb_test, detail_file)
+ err, siz = test_dipole(dp, data, numb_test, detail_file, atomic)
elif dp.model_type == "polar":
err, siz = test_polar(dp, data, numb_test, detail_file, global_polar=False)
elif dp.model_type == "global_polar":
@@ -550,6 +550,7 @@ def test_dipole(
data: DeepmdData,
numb_test: int,
detail_file: Optional[str],
+ has_atom_dipole: bool,
) -> Tuple[List[np.ndarray], List[int]]:
"""Test energy type model.
@@ -563,6 +564,8 @@ def test_dipole(
munber of tests to do
detail_file : Optional[str]
file where test details will be output
+ has_atom_dipole : bool
+ whether atomic dipole is provided
Returns
-------
@@ -570,12 +573,22 @@ def test_dipole(
arrays with results and their shapes
"""
data.add(
- "dipole", 3, atomic=True, must=True, high_prec=False, type_sel=dp.get_sel_type()
+ "dipole", 3, atomic=has_atom_dipole, must=True, high_prec=False, type_sel=dp.get_sel_type()
)
test_data = data.get_test()
dipole, numb_test, _ = run_test(dp, test_data, numb_test)
+
+ # do summation in atom dimension
+ if has_atom_dipole == False:
+ dipole = np.reshape(dipole,(dipole.shape[0], -1, 3))
+ atoms = dipole.shape[1]
+ dipole = np.sum(dipole,axis=1)
+
l2f = l2err(dipole - test_data["dipole"][:numb_test])
+ if has_atom_dipole == False:
+ l2f = l2f / atoms
+
log.info(f"# number of test data : {numb_test:d}")
log.info(f"Dipole RMSE : {l2f:e} eV/A")
diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py
index 685815db17..790cc227f6 100755
--- a/deepmd/entrypoints/train.py
+++ b/deepmd/entrypoints/train.py
@@ -253,9 +253,9 @@ def _do_work(jdata: Dict[str, Any], run_opt: RunOptions):
# get batch sizes
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")
+ stop_batch = j_must_have(jdata["training"], "numb_steps")
sys_probs = jdata["training"].get("sys_probs")
- auto_prob_style = jdata["training"].get("auto_prob_style", "prob_sys_size")
+ auto_prob = jdata["training"].get("auto_prob", "prob_sys_size")
# setup data modifier
modifier: Optional[DipoleChargeModifier]
@@ -283,8 +283,9 @@ def _do_work(jdata: Dict[str, Any], run_opt: RunOptions):
set_prefix=set_pfx,
type_map=ipt_type_map,
modifier=modifier,
+ trn_all_set = True
)
- data.print_summary(run_opt, sys_probs=sys_probs, auto_prob_style=auto_prob_style)
+ data.print_summary(run_opt, sys_probs=sys_probs, auto_prob_style=auto_prob)
data.add_dict(data_requirement)
# build the model with stats from the first system
diff --git a/deepmd/loss/tensor.py b/deepmd/loss/tensor.py
index ceff94e14b..addccdcadf 100644
--- a/deepmd/loss/tensor.py
+++ b/deepmd/loss/tensor.py
@@ -12,24 +12,62 @@ class TensorLoss () :
def __init__ (self, jdata, **kwarg) :
try:
model = kwarg['model']
- type_sel = model.get_sel_type()
+ self.type_sel = model.get_sel_type()
except :
- type_sel = None
+ self.type_sel = None
self.tensor_name = kwarg['tensor_name']
self.tensor_size = kwarg['tensor_size']
self.label_name = kwarg['label_name']
- self.atomic = kwarg.get('atomic', True)
+ self.atomic = kwarg.get('atomic', None)
if jdata is not None:
self.scale = jdata.get('scale', 1.0)
else:
self.scale = 1.0
+
+ # YHT: added for global / local dipole combination
+ if self.atomic is True: # upper regulation, will control the lower behavior
+ self.local_weight,self.global_weight = 1.0,0.0
+ elif self.atomic is False: # upper regulation, will control the lower behavior
+ self.local_weight,self.global_weight = 0.0,1.0
+ else: # self.atomic is None, let the loss parameter decide which mode to use
+ if jdata is not None:
+ self.local_weight = jdata.get('pref_atomic_' + self.tensor_name,None)
+ self.global_weight = jdata.get('pref_' + self.tensor_name,None)
+
+ # get the input parameter first
+ if self.local_weight is None and self.global_weight is None:
+ # default: downward compatibility, using local mode
+ self.local_weight , self.global_weight = 1.0, 0.0
+ self.atomic = True
+ elif self.local_weight is None and self.global_weight is not None:
+ # using global mode only, normalize to 1
+ assert self.global_weight > 0.0, AssertionError('assign a zero weight to global dipole without setting a local weight')
+ self.local_weight = 0.0
+ self.atomic = False
+ elif self.local_weight is not None and self.global_weight is None:
+ assert self.local_weight > 0.0, AssertionError('assign a zero weight to local dipole without setting a global weight')
+ self.global_weight = 0.0
+ self.atomic = True
+ else: # Both are not None
+ self.atomic = True if self.local_weight != 0.0 else False
+ assert (self.local_weight >0.0) or (self.global_weight>0.0), AssertionError('can not assian zero weight to both local and global mode')
+
+ # normalize, not do according to Han Wang's suggestion
+ #temp_sum = self.local_weight + self.global_weight
+ #self.local_weight /= temp_sum
+ #self.global_weight /= temp_sum
+
+ else: # Nothing been set, use default setting
+ self.local_weight,self.global_weight = 1.0,0.0
+ self.atomic = True
+
# data required
add_data_requirement(self.label_name,
self.tensor_size,
atomic=self.atomic,
must=True,
high_prec=False,
- type_sel = type_sel)
+ type_sel = self.type_sel)
def build (self,
learning_rate,
@@ -39,22 +77,67 @@ def build (self,
suffix):
polar_hat = label_dict[self.label_name]
polar = model_dict[self.tensor_name]
- l2_loss = tf.reduce_mean( tf.square(self.scale*(polar - polar_hat)), name='l2_'+suffix)
- more_loss = {'nonorm': l2_loss}
- if not self.atomic :
- atom_norm = 1./ global_cvt_2_tf_float(natoms[0])
- l2_loss = l2_loss * atom_norm
+
+ # YHT: added for global / local dipole combination
+ l2_loss = global_cvt_2_tf_float(0.0)
+ more_loss = {
+ "local_loss":global_cvt_2_tf_float(0.0),
+ "global_loss":global_cvt_2_tf_float(0.0)
+ }
+
+
+ if self.local_weight > 0.0:
+ local_loss = tf.reduce_mean( tf.square(self.scale*(polar - polar_hat)), name='l2_'+suffix)
+ more_loss['local_loss'] = local_loss
+ l2_loss += self.local_weight * local_loss
+ self.l2_loss_local_summary = tf.summary.scalar('l2_local_loss',
+ tf.sqrt(more_loss['local_loss']))
+
+
+ if self.global_weight > 0.0: # Need global loss
+ atoms = 0
+ if self.type_sel is not None:
+ for w in self.type_sel:
+ atoms += natoms[2+w]
+ else:
+ atoms = natoms[0]
+ nframes = tf.shape(polar)[0] // self.tensor_size // atoms
+ # get global results
+ global_polar = tf.reshape(tf.reduce_sum(tf.reshape(
+ polar, [nframes, -1, self.tensor_size]), axis=1),[-1])
+ if self.atomic: # If label is local, however
+ global_polar_hat = tf.reshape(tf.reduce_sum(tf.reshape(
+ polar_hat, [nframes, -1, self.tensor_size]), axis=1),[-1])
+ else:
+ global_polar_hat = polar_hat
+
+ global_loss = tf.reduce_mean( tf.square(self.scale*(global_polar - global_polar_hat)), name='l2_'+suffix)
+
+ more_loss['global_loss'] = global_loss
+ self.l2_loss_global_summary = tf.summary.scalar('l2_global_loss',
+ tf.sqrt(more_loss['global_loss']) / global_cvt_2_tf_float(atoms))
+
+ # YHT: should only consider atoms with dipole, i.e. atoms
+ # atom_norm = 1./ global_cvt_2_tf_float(natoms[0])
+ atom_norm = 1./ global_cvt_2_tf_float(atoms)
+ global_loss *= atom_norm
+
+ l2_loss += self.global_weight * global_loss
+
+ self.l2_more = more_loss
self.l2_l = l2_loss
- self.l2_more = more_loss['nonorm']
self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
return l2_loss, more_loss
- @staticmethod
- def print_header():
+ def print_header(self):
prop_fmt = ' %11s %11s'
print_str = ''
print_str += prop_fmt % ('rmse_tst', 'rmse_trn')
+ if self.local_weight > 0.0:
+ print_str += prop_fmt % ('rmse_lc_tst', 'rmse_lc_trn')
+ if self.global_weight > 0.0:
+ print_str += prop_fmt % ('rmse_gl_tst', 'rmse_gl_trn')
return print_str
def print_on_training(self,
@@ -65,7 +148,20 @@ def print_on_training(self,
feed_dict_test,
feed_dict_batch) :
- run_data = [self.l2_l]
+ # YHT: added to calculate the atoms number
+ atoms = 0
+ if self.type_sel is not None:
+ for w in self.type_sel:
+ atoms += natoms[2+w]
+ else:
+ atoms = natoms[0]
+
+ run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']]
+ summary_list = [self.l2_loss_summary]
+ if self.local_weight > 0.0:
+ summary_list.append(self.l2_loss_local_summary)
+ if self.global_weight > 0.0:
+ summary_list.append(self.l2_loss_global_summary)
# first train data
error_train = sess.run(run_data, feed_dict=feed_dict_batch)
@@ -73,7 +169,8 @@ def print_on_training(self,
# than test data, if tensorboard log writter is present, commpute summary
# and write tensorboard logs
if tb_writer:
- summary_merged_op = tf.summary.merge([self.l2_loss_summary])
+ #summary_merged_op = tf.summary.merge([self.l2_loss_summary])
+ summary_merged_op = tf.summary.merge(summary_list)
run_data.insert(0, summary_merged_op)
test_out = sess.run(run_data, feed_dict=feed_dict_test)
@@ -82,10 +179,14 @@ def print_on_training(self,
summary = test_out.pop(0)
tb_writer.add_summary(summary, cur_batch)
- error_test = test_out[0]
+ error_test = test_out
print_str = ""
prop_fmt = " %11.2e %11.2e"
- print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
+ print_str += prop_fmt % (np.sqrt(error_test[0]), np.sqrt(error_train[0]))
+ if self.local_weight > 0.0:
+ print_str += prop_fmt % (np.sqrt(error_test[1]), np.sqrt(error_train[1]) )
+ if self.global_weight > 0.0:
+ print_str += prop_fmt % (np.sqrt(error_test[2])/atoms, np.sqrt(error_train[2])/atoms)
return print_str
diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py
index 1ea177fd95..1ae50e1e89 100644
--- a/deepmd/train/trainer.py
+++ b/deepmd/train/trainer.py
@@ -251,7 +251,7 @@ def _init_param(self, jdata):
.add('tensorboard', bool, default = False)\
.add('tensorboard_log_dir',str, default = 'log')\
.add('sys_probs', list )\
- .add('auto_prob_style', str, default = "prob_sys_size")
+ .add('auto_prob', str, default = "prob_sys_size")
tr_data = tr_args.parse(training_param)
# not needed
# self.numb_test = tr_data['numb_test']
@@ -266,7 +266,7 @@ def _init_param(self, jdata):
self.tensorboard = tr_data['tensorboard']
self.tensorboard_log_dir = tr_data['tensorboard_log_dir']
self.sys_probs = tr_data['sys_probs']
- self.auto_prob_style = tr_data['auto_prob_style']
+ self.auto_prob_style = tr_data['auto_prob']
self.useBN = False
if fitting_type == 'ener' and self.fitting.get_numb_fparam() > 0 :
self.numb_fparam = self.fitting.get_numb_fparam()
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index ac993ef107..310090adfd 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -299,7 +299,7 @@ def model_args ():
return ca
-def learning_rate_args():
+def learning_rate_exp():
doc_start_lr = 'The learning rate the start of the training.'
doc_stop_lr = 'The desired learning rate at the end of the training.'
doc_decay_steps = 'The learning rate is decaying every this number of training steps.'
@@ -309,9 +309,24 @@ def learning_rate_args():
Argument("stop_lr", float, optional = True, default = 1e-8, doc = doc_stop_lr),
Argument("decay_steps", int, optional = True, default = 5000, doc = doc_decay_steps)
]
+ return args
+
+
+def learning_rate_variant_type_args():
+ doc_lr = 'The type of the learning rate. Current type `exp`, the exponentially decaying learning rate is supported.'
+
+ return Variant("type",
+ [Argument("exp", dict, learning_rate_exp())],
+ optional = True,
+ default_tag = 'exp',
+ doc = doc_lr)
- doc_lr = "The learning rate options"
- return Argument("learning_rate", dict, args, [], doc = doc_lr)
+
+def learning_rate_args():
+ doc_lr = "The definitio of learning rate"
+ return Argument("learning_rate", dict, [],
+ [learning_rate_variant_type_args()],
+ doc = doc_lr)
def start_pref(item):
@@ -361,15 +376,16 @@ def loss_args():
return ca
def training_args():
+ link_sys = make_link("systems", "training/systems")
doc_systems = 'The data systems. This key can be provided with a listthat specifies the systems, or be provided with a string by which the prefix of all systems are given and the list of the systems is automatically generated.'
- doc_set_prefix = 'The prefix of the sets in the systems.'
+ doc_set_prefix = f'The prefix of the sets in the {link_sys}.'
doc_stop_batch = 'Number of training batch. Each training uses one batch of data.'
- doc_batch_size = 'This key can be \n\n\
-- list: the length of which is the same as the `systems`. The batch size of each system is given by the elements of the list.\n\n\
-- int: all `systems` uses the same batch size.\n\n\
-- string "auto": automatically determines the batch size os that the batch_size times the number of atoms in the system is no less than 32.\n\n\
-- string "auto:N": automatically determines the batch size os that the batch_size times the number of atoms in the system is no less than N.'
- doc_seed = 'The random seed for training.'
+ doc_batch_size = f'This key can be \n\n\
+- list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\
+- int: all {link_sys} use the same batch size.\n\n\
+- string "auto": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than 32.\n\n\
+- string "auto:N": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than N.'
+ doc_seed = 'The random seed for getting frames from the training data set.'
doc_disp_file = 'The file for printing learning curve.'
doc_disp_freq = 'The frequency of printing learning curve.'
doc_numb_test = 'Number of frames used for the test during training.'
@@ -379,12 +395,21 @@ def training_args():
doc_time_training = 'Timing durining training.'
doc_profiling = 'Profiling during training.'
doc_profiling_file = 'Output file for profiling.'
+ doc_train_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\
+- "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\
+- "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\
+- "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is devided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.'
+ doc_train_sys_probs = "A list of float, should be of the same length as `train_systems`, specifying the probability of each system."
+ doc_tensorboard = 'Enable tensorboard'
+ doc_tensorboard_log_dir = 'The log directory of tensorboard outputs'
args = [
- Argument("systems", [list,str], optional = False, doc = doc_systems),
+ Argument("systems", [list,str], optional = False, doc = doc_systems, alias = ["trn_systems"]),
Argument("set_prefix", str, optional = True, default = 'set', doc = doc_set_prefix),
- Argument("stop_batch", int, optional = False, doc = doc_stop_batch),
- Argument("batch_size", [list,int,str], optional = True, default = 'auto', doc = doc_batch_size),
+ Argument("auto_prob", str, optional = True, default = "prob_sys_size", doc = doc_train_auto_prob_style, alias = ["trn_auto_prob", "auto_prob_style"]),
+ Argument("sys_probs", list, optional = True, default = None, doc = doc_train_sys_probs, alias = ["trn_sys_probs"]),
+ Argument("batch_size", [list,int,str], optional = True, default = 'auto', doc = doc_batch_size, alias = ["trn_batch_size"]),
+ Argument("numb_steps", int, optional = False, doc = doc_stop_batch, alias = ["stop_batch"]),
Argument("seed", [int,None], optional = True, doc = doc_seed),
Argument("disp_file", str, optional = True, default = 'lcueve.out', doc = doc_disp_file),
Argument("disp_freq", int, optional = True, default = 1000, doc = doc_disp_freq),
@@ -394,7 +419,9 @@ def training_args():
Argument("disp_training", bool, optional = True, default = True, doc = doc_disp_training),
Argument("time_training", bool, optional = True, default = True, doc = doc_time_training),
Argument("profiling", bool, optional = True, default = False, doc = doc_profiling),
- Argument("profiling_file", str, optional = True, default = 'timeline.json', doc = doc_profiling_file)
+ Argument("profiling_file", str, optional = True, default = 'timeline.json', doc = doc_profiling_file),
+ Argument("tensorboard", bool, optional = True, default = False, doc = doc_tensorboard),
+ Argument("tensorboard_log_dir", str, optional = True, default = 'log', doc = doc_tensorboard_log_dir),
]
doc_training = 'The training options'
@@ -435,7 +462,7 @@ def normalize(data):
base = Argument("base", dict, [ma, lra, la, ta])
data = base.normalize_value(data, trim_pattern = "_*")
- base.check_value(data)
+ base.check_value(data, strict = True)
return data
diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py
index 4b8edebb2d..b52c6e17c8 100644
--- a/deepmd/utils/data.py
+++ b/deepmd/utils/data.py
@@ -19,7 +19,8 @@ def __init__ (self,
set_prefix : str = 'set',
shuffle_test : bool = True,
type_map : List[str] = None,
- modifier = None) :
+ modifier = None,
+ trn_all_set : bool = False) :
"""
Constructor
@@ -35,6 +36,8 @@ def __init__ (self,
Gives the name of different atom types
modifier
Data modifier that has the method `modify_data`
+ trn_all_set
+ Use all sets as training dataset. Otherwise, if the number of sets is more than 1, the last set is left for test.
"""
self.dirs = glob.glob (os.path.join(sys_path, set_prefix + ".*"))
self.dirs.sort()
@@ -57,10 +60,13 @@ def __init__ (self,
self.idx_map = self._make_idx_map(self.atom_type)
# train dirs
self.test_dir = self.dirs[-1]
- if len(self.dirs) == 1 :
+ if trn_all_set:
self.train_dirs = self.dirs
- else :
- self.train_dirs = self.dirs[:-1]
+ else:
+ if len(self.dirs) == 1 :
+ self.train_dirs = self.dirs
+ else :
+ self.train_dirs = self.dirs[:-1]
self.data_dict = {}
# add box and coord
self.add('box', 9, must = True)
diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py
index 557a2c6bd0..188366540b 100644
--- a/deepmd/utils/data_system.py
+++ b/deepmd/utils/data_system.py
@@ -26,7 +26,8 @@ def __init__ (self,
set_prefix : str = 'set',
shuffle_test : bool = True,
type_map : List[str] = None,
- modifier = None) :
+ modifier = None,
+ trn_all_set = False) :
"""
Constructor
@@ -48,6 +49,8 @@ def __init__ (self,
Gives the name of different atom types
modifier
Data modifier that has the method `modify_data`
+ trn_all_set
+ Use all sets as training dataset. Otherwise, if the number of sets is more than 1, the last set is left for test.
"""
# init data
self.rcut = rcut
@@ -55,11 +58,15 @@ def __init__ (self,
self.nsystems = len(self.system_dirs)
self.data_systems = []
for ii in self.system_dirs :
- self.data_systems.append(DeepmdData(ii,
- set_prefix=set_prefix,
- shuffle_test=shuffle_test,
- type_map = type_map,
- modifier = modifier))
+ self.data_systems.append(
+ DeepmdData(
+ ii,
+ set_prefix=set_prefix,
+ shuffle_test=shuffle_test,
+ type_map = type_map,
+ modifier = modifier,
+ trn_all_set = trn_all_set
+ ))
# batch size
self.batch_size = batch_size
if isinstance(self.batch_size, int) :
@@ -260,7 +267,8 @@ def _get_sys_probs(self,
auto_prob_style) :
if sys_probs is None :
if auto_prob_style == "prob_uniform" :
- prob = None
+ prob_v = 1./float(self.nsystems)
+ prob = [prob_v for ii in range(self.nsystems)]
elif auto_prob_style == "prob_sys_size" :
prob = self.prob_nbatches
elif auto_prob_style[:14] == "prob_sys_size;" :
@@ -273,9 +281,9 @@ def _get_sys_probs(self,
def get_batch (self,
- sys_idx = None,
- sys_probs = None,
- auto_prob_style = "prob_sys_size") :
+ sys_idx : int = None,
+ sys_probs : List[float] = None,
+ auto_prob_style : str = "prob_sys_size") :
"""
Get a batch of data from the data systems
@@ -289,7 +297,7 @@ def get_batch (self,
The probabilitis of systems to get the batch.
Summation of positive elements of this list should be no greater than 1.
Element of this list can be negative, the probability of the corresponding system is determined automatically by the number of batches in the system.
- auto_prob_style: float
+ auto_prob_style: str
Determine the probability of systems automatically. The method is assigned by this key and can be
- "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()
- "prob_sys_size" : the probability of a system is proportional to the number of batches in the system
@@ -459,9 +467,12 @@ def _process_sys_probs(self, sys_probs) :
assigned_sum_prob = np.sum(type_filter * sys_probs)
assert assigned_sum_prob <= 1, "the sum of assigned probability should be less than 1"
rest_sum_prob = 1. - assigned_sum_prob
- rest_nbatch = (1 - type_filter) * self.nbatches
- rest_prob = rest_sum_prob * rest_nbatch / np.sum(rest_nbatch)
- ret_prob = rest_prob + type_filter * sys_probs
+ if rest_sum_prob != 0 :
+ rest_nbatch = (1 - type_filter) * self.nbatches
+ rest_prob = rest_sum_prob * rest_nbatch / np.sum(rest_nbatch)
+ ret_prob = rest_prob + type_filter * sys_probs
+ else :
+ ret_prob = sys_probs
assert np.sum(ret_prob) == 1, "sum of probs should be 1"
return ret_prob
diff --git a/doc/train-input-auto.rst b/doc/train-input-auto.rst
index aa85750cbc..6da37d5264 100644
--- a/doc/train-input-auto.rst
+++ b/doc/train-input-auto.rst
@@ -23,6 +23,15 @@ model:
The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics.
+ .. raw:: html
+
+
+ data_stat_protect:
+ | type: ``float``, optional, default: ``0.01``
+ | argument path: ``model/data_stat_protect``
+
+ Protect parameter for atomic energy regression.
+
.. raw:: html
@@ -233,7 +242,7 @@ model:
| type: ``str``, optional, default: ``float64``
| argument path: ``model/descriptor[se_a]/precision``
- The precision of the embedding net parameters, supported options are "float64", "float32", "float16".
+ The precision of the embedding net parameters, supported options are "default", "float16", "float32", "float64".
.. raw:: html
@@ -347,7 +356,7 @@ model:
| type: ``str``, optional, default: ``float64``
| argument path: ``model/descriptor[se_r]/precision``
- The precision of the embedding net parameters, supported options are "float64", "float32", "float16".
+ The precision of the embedding net parameters, supported options are "default", "float16", "float32", "float64".
.. raw:: html
@@ -452,7 +461,7 @@ model:
| type: ``str``, optional, default: ``float64``
| argument path: ``model/descriptor[se_a_3be]/precision``
- The precision of the embedding net parameters, supported options are "float64", "float32", "float16".
+ The precision of the embedding net parameters, supported options are "default", "float16", "float32", "float64".
.. raw:: html
@@ -472,15 +481,6 @@ model:
Random seed for parameter initialization
- .. raw:: html
-
-
- exclude_types:
- | type: ``list``, optional, default: ``[]``
- | argument path: ``model/descriptor[se_a_3be]/exclude_types``
-
- The Excluded types
-
.. raw:: html
@@ -575,7 +575,7 @@ model:
| type: ``str``, optional, default: ``float64``
| argument path: ``model/descriptor[se_a_tpe]/precision``
- The precision of the embedding net parameters, supported options are "float64", "float32", "float16".
+ The precision of the embedding net parameters, supported options are "default", "float16", "float32", "float64".
.. raw:: html
@@ -757,7 +757,7 @@ model:
| type: ``str``, optional, default: ``float64``
| argument path: ``model/fitting_net[ener]/precision``
- The precision of the fitting net parameters, supported options are "float64", "float32", "float16".
+ The precision of the fitting net parameters, supported options are "default", "float16", "float32", "float64".
.. raw:: html
@@ -848,7 +848,7 @@ model:
| type: ``str``, optional, default: ``float64``
| argument path: ``model/fitting_net[dipole]/precision``
- The precision of the fitting net parameters, supported options are "float64", "float32", "float16".
+ The precision of the fitting net parameters, supported options are "default", "float16", "float32", "float64".
.. raw:: html
@@ -908,7 +908,7 @@ model:
| type: ``str``, optional, default: ``float64``
| argument path: ``model/fitting_net[polar]/precision``
- The precision of the fitting net parameters, supported options are "float64", "float32", "float16".
+ The precision of the fitting net parameters, supported options are "default", "float16", "float32", "float64".
.. raw:: html
@@ -995,7 +995,7 @@ model:
| type: ``str``, optional, default: ``float64``
| argument path: ``model/fitting_net[global_polar]/precision``
- The precision of the fitting net parameters, supported options are "float64", "float32", "float16".
+ The precision of the fitting net parameters, supported options are "default", "float16", "float32", "float64".
.. raw:: html
@@ -1047,7 +1047,7 @@ model:
loss:
- | type: ``dict``
+ | type: ``dict``, optional
| argument path: ``loss``
The definition of loss function. The type of the loss depends on the type of the fitting. For fitting type `ener`, the prefactors before energy, force, virial and atomic energy losses may be provided. For fitting type `dipole`, `polar` and `global_polar`, the loss may be an empty `dict` or unset.
@@ -1160,32 +1160,50 @@ learning_rate:
| type: ``dict``
| argument path: ``learning_rate``
- The learning rate options
+ The definitio of learning rate
+
+
+ Depending on the value of *type*, different sub args are accepted.
+
+ .. raw:: html
+
+
+ type:
+ | type: ``str`` (flag key), default: ``exp``
+ | argument path: ``learning_rate/type``
+
+ The type of the learning rate. Current type `exp`, the exponentially decaying learning rate is supported.
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``exp``:
.. raw:: html
-
+
start_lr:
| type: ``float``, optional, default: ``0.001``
- | argument path: ``learning_rate/start_lr``
+ | argument path: ``learning_rate[exp]/start_lr``
The learning rate the start of the training.
.. raw:: html
-
+
stop_lr:
| type: ``float``, optional, default: ``1e-08``
- | argument path: ``learning_rate/stop_lr``
+ | argument path: ``learning_rate[exp]/stop_lr``
The desired learning rate at the end of the training.
.. raw:: html
-
+
decay_steps:
| type: ``int``, optional, default: ``5000``
- | argument path: ``learning_rate/decay_steps``
+ | argument path: ``learning_rate[exp]/decay_steps``
The learning rate is decaying every this number of training steps.
@@ -1215,16 +1233,31 @@ training:
| type: ``str``, optional, default: ``set``
| argument path: ``training/set_prefix``
- The prefix of the sets in the systems.
+ The prefix of the sets in the `systems <#training/systems>`__.
.. raw:: html
-
- stop_batch:
- | type: ``int``
- | argument path: ``training/stop_batch``
+
+ auto_prob:
+ | type: ``str``, optional, default: ``prob_sys_size``
+ | argument path: ``training/auto_prob``
- Number of training batch. Each training uses one batch of data.
+ Determine the probability of systems automatically. The method is assigned by this key and can be
+
+ - "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()
+
+ - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system
+
+ - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is devided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.
+
+ .. raw:: html
+
+
+ sys_probs:
+ | type: ``NoneType`` | ``list``, optional, default: ``None``
+ | argument path: ``training/sys_probs``
+
+ A list of float, should be of the same length as `train_systems`, specifying the probability of each system.
.. raw:: html
@@ -1235,13 +1268,22 @@ training:
This key can be
- - list: the length of which is the same as the `systems`. The batch size of each system is given by the elements of the list.
+ - list: the length of which is the same as the `systems <#training/systems>`__. The batch size of each system is given by the elements of the list.
+
+ - int: all `systems <#training/systems>`__ use the same batch size.
+
+ - string "auto": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than 32.
+
+ - string "auto:N": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than N.
- - int: all `systems` uses the same batch size.
+ .. raw:: html
- - string "auto": automatically determines the batch size os that the batch_size times the number of atoms in the system is no less than 32.
+
+ numb_steps:
+ | type: ``int``
+ | argument path: ``training/numb_steps``
- - string "auto:N": automatically determines the batch size os that the batch_size times the number of atoms in the system is no less than N.
+ Number of training batch. Each training uses one batch of data.
.. raw:: html
@@ -1250,7 +1292,7 @@ training:
| type: ``int`` | ``NoneType``, optional
| argument path: ``training/seed``
- The random seed for training.
+ The random seed for getting frames from the training data set.
.. raw:: html
@@ -1333,3 +1375,21 @@ training:
Output file for profiling.
+ .. raw:: html
+
+
+ tensorboard:
+ | type: ``bool``, optional, default: ``False``
+ | argument path: ``training/tensorboard``
+
+ Enable tensorboard
+
+ .. raw:: html
+
+
+ tensorboard_log_dir:
+ | type: ``str``, optional, default: ``log``
+ | argument path: ``training/tensorboard_log_dir``
+
+ The log directory of tensorboard outputs
+
diff --git a/requirements.txt b/requirements.txt
index 3b80afd9c0..c8a81216ab 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
numpy
scipy
pyyaml
-dargs
+dargs >= 0.2.0
tqdm
-typing_extensions
\ No newline at end of file
+typing_extensions
diff --git a/source/tests/test_argument_parser.py b/source/tests/test_argument_parser.py
index 356dd430c1..a7e42ad81e 100644
--- a/source/tests/test_argument_parser.py
+++ b/source/tests/test_argument_parser.py
@@ -262,7 +262,7 @@ def test_parser_test(self):
"--numb-test": dict(type=int, value=1),
"--rand-seed": dict(type=(int, type(None)), value=12321),
"--detail-file": dict(type=(str, type(None)), value="TARGET.FILE"),
- "--atomic-energy": dict(type=bool),
+ "--atomic": dict(type=bool),
}
self.run_test(command="test", mapping=ARGS)