required=True, help="Which mode to use train or retrain") args = vars(parser.parse_args()) # args = {"filepath":"E:/Benutzer/Patrick/PostDoc/Projects ML/NeuralNet4/NNfit0/nac_0",'index' : 0,"gpus":0} fstdout = open( os.path.join(args['filepath'], "fitlog_" + str(args['index']) + ".txt"), 'w') sys.stderr = fstdout sys.stdout = fstdout print("Input argpars:", args) from NNsMD.nn_pes_src.device import set_gpu set_gpu([int(args['gpus'])]) print("Logic Devices:", tf.config.experimental.list_logical_devices('GPU')) from NNsMD.utils.callbacks import EarlyStopping, lr_lin_reduction, lr_exp_reduction, lr_step_reduction from NNsMD.models.mlp_nac2 import NACModel2 from NNsMD.datasets.general import load_hyp from NNsMD.datasets.general import split_validation_training_index from NNsMD.scaler.nac import NACStandardScaler from NNsMD.scaler.general import SegmentStandardScaler from NNsMD.utils.loss import ScaledMeanAbsoluteError, get_lr_metric, r2_metric, NACphaselessLoss from NNsMD.plots.loss import plot_loss_curves, plot_learning_curve from NNsMD.plots.pred import plot_scatter_prediction from NNsMD.plots.error import plot_error_vec_mean, plot_error_vec_max def train_model_nac(i=0, outdir=None, mode='training'):
def __init__(self, keywords=None, id=None): set_gpu([]) #No GPU for prediction title = keywords['control']['title'] variables = keywords['nn'].copy() modeldir = variables['modeldir'] data = variables['data'] nn_eg_type = variables['nn_eg_type'] nn_nac_type = variables['nn_nac_type'] nn_soc_type = variables['nn_soc_type'] hyp_eg = variables['eg'].copy() hyp_nac = variables['nac'].copy() hyp_eg2 = variables['eg2'].copy() hyp_nac2 = variables['nac2'].copy() hyp_soc = variables['soc'].copy() hyp_soc2 = variables['soc2'].copy() eg_unit = variables['eg_unit'] nac_unit = variables['nac_unit'] soc_unit = variables['soc_unit'] seed = variables['ml_seed'] permute = variables['permute_map'] gpu = variables['gpu'] self.jobtype = keywords['control']['jobtype'] self.version = keywords['version'] self.ncpu = keywords['control']['ml_ncpu'] self.pred_data = variables['pred_data'] self.train_mode = variables['train_mode'] self.shuffle = variables['shuffle'] self.natom = data.natom self.nstate = data.nstate self.nnac = data.nnac self.nsoc = data.nsoc ## set hyperparamters hyp_dict_eg = SetHyperEG(hyp_eg, eg_unit, data.info) hyp_dict_eg2 = SetHyperEG(hyp_eg2, eg_unit, data.info) hyp_dict_nac = SetHyperNAC(hyp_nac, nac_unit, data.info) hyp_dict_nac2 = SetHyperNAC(hyp_nac2, nac_unit, data.info) hyp_dict_soc = SetHyperSOC(hyp_soc, soc_unit, data.info) hyp_dict_soc2 = SetHyperSOC(hyp_soc2, soc_unit, data.info) ## retraining has some bug at the moment, do not use if self.train_mode not in ['training', 'retraining', 'resample']: self.train_mode = 'training' if id == None or id == 1: self.name = f"NN-{title}" else: self.name = f"NN-{title}-{id}" self.silent = variables['silent'] self.x = data.x self.pred_x = data.pred_x self.pred_y = data.pred_y ## convert unit of energy and force. au or si. data are in au. if eg_unit == 'si': self.H_to_eV = 27.211396132 self.H_Bohr_to_eV_A = 27.211396132 / 0.529177249 self.keep_eV = 1 self.keep_eVA = 1 else: self.H_to_eV = 1 self.H_Bohr_to_eV_A = 1 self.keep_eV = 27.211396132 self.keep_eVA = 27.211396132 / 0.529177249 if nac_unit == 'si': self.Bohr_to_A = 0.529177249 / 27.211396132 # convert to eV/A self.keep_A = 1 elif nac_unit == 'au': self.Bohr_to_A = 1 # convert to Eh/B self.keep_A = 0.529177249 / 27.211396132 elif nac_unit == 'eha': self.Bohr_to_A = 0.529177249 # convert to Eh/A self.keep_A = 1 / 27.211396132 else: self.Bohr_to_A = 1 # convert to Eh/B self.keep_A = 0.529177249 / 27.211396132 ## combine y_dict self.y_dict = {} if nn_eg_type > 0: self.y_dict['energy_gradient'] = [ data.energy * self.H_to_eV, data.grad * self.H_Bohr_to_eV_A ] if nn_nac_type > 0: self.y_dict['nac'] = data.nac / self.Bohr_to_A if nn_soc_type > 0: self.y_dict['soc'] = data.soc ## check permuation map self.x, self.y_dict = PermuteMap(self.x, self.y_dict, permute, hyp_eg['val_split']) ## combine hypers self.hyper = {} if nn_eg_type == 1: # same architecture with different weight self.hyper['energy_gradient'] = hyp_dict_eg elif nn_eg_type > 1: self.hyper['energy_gradient'] = [hyp_dict_eg, hyp_dict_eg2] if nn_nac_type == 1: # same architecture with different weight self.hyper['nac'] = hyp_dict_nac elif nn_nac_type > 1: self.hyper['nac'] = [hyp_dict_nac, hyp_dict_nac2] if nn_soc_type == 1: # same architecture with different weight self.hyper['soc'] = hyp_dict_soc elif nn_soc_type > 1: self.hyper['soc'] = [hyp_dict_soc, hyp_dict_soc2] ## setup GPU list self.gpu_list = {} if gpu == 1: self.gpu_list['energy_gradient'] = [0, 0] self.gpu_list['nac'] = [0, 0] self.gpu_list['soc'] = [0, 0] elif gpu == 2: self.gpu_list['energy_gradient'] = [0, 1] self.gpu_list['nac'] = [0, 1] self.gpu_list['soc'] = [0, 1] elif gpu == 3: self.gpu_list['energy_gradient'] = [0, 0] self.gpu_list['nac'] = [1, 1] self.gpu_list['soc'] = [2, 2] elif gpu == 4: self.gpu_list['energy_gradient'] = [0, 1] self.gpu_list['nac'] = [2, 2] self.gpu_list['soc'] = [3, 3] elif gpu == 5: self.gpu_list['energy_gradient'] = [0, 1] self.gpu_list['nac'] = [2, 3] self.gpu_list['soc'] = [4, 4] elif gpu == 6: self.gpu_list['energy_gradient'] = [0, 1] self.gpu_list['nac'] = [2, 3] self.gpu_list['soc'] = [4, 5] ## initialize model if modeldir == None or id not in [None, 1]: self.model = NeuralNetPes(self.name) else: self.model = NeuralNetPes(modeldir)