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)