scale tensor loss

Author: Han Wang

source/train/Fitting.py

@@ -395,6 +395,12 @@ def __init__ (self, jdata, descrpt) :
         self.seed = class_data['seed']
         self.diag_shift = class_data['diag_shift']
         self.scale = class_data['scale']
+        if type(self.sel_type) is not list:
+            self.sel_type = [self.sel_type]
+        if type(self.diag_shift) is not list:
+            self.diag_shift = [self.diag_shift]
+        if type(self.scale) is not list:
+            self.scale = [self.scale]
         self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
         self.dim_rot_mat = self.dim_rot_mat_1 * 3
         self.useBN = False

source/train/Loss.py

@@ -284,6 +284,10 @@ def __init__ (self, jdata, **kwarg) :
         self.tensor_size = kwarg['tensor_size']
         self.label_name = kwarg['label_name']
         self.atomic = kwarg.get('atomic', True)
+        if jdata is not None:
+            self.scale = jdata.get('scale', 1.0)
+        else:
+            self.scale = 1.0
         # data required
         add_data_requirement(self.label_name, 
                              self.tensor_size, 
@@ -300,7 +304,7 @@ def build (self,
                suffix):        
         polar_hat = label_dict[self.label_name]
         polar = model_dict[self.tensor_name]
-        l2_loss = tf.reduce_mean( tf.square(polar - polar_hat), name='l2_'+suffix)
+        l2_loss = tf.reduce_mean( tf.square(self.scale*(polar - polar_hat)), name='l2_'+suffix)
         if not self.atomic :
             atom_norm  = 1./ global_cvt_2_tf_float(natoms[0]) 
             l2_loss = l2_loss * atom_norm