diff --git a/simple_nn_v2/models/data_handler.py b/simple_nn_v2/models/data_handler.py index c877fa9..6c98223 100644 --- a/simple_nn_v2/models/data_handler.py +++ b/simple_nn_v2/models/data_handler.py @@ -63,7 +63,7 @@ def save_filename(self): for f in self.filelist: tmp_dict = torch.load(f) tmp_dict['filename'] = f - torch.save(tmp_dict,f) + torch.save(tmp_dict, f) #Used in Structure rmse class StructlistDataset(FilelistDataset): @@ -182,7 +182,7 @@ def my_collate(batch, atom_types, scale_factor=None, pca=None, pca_min_whiten_le if use_stress: S = torch.cat(S, axis=0) - return {'x': x, 'dx': dx, 'da': da, 'n': n, 'E': E, 'F': F, 'S': S, 'sp_idx': sparse_index, 'struct_weight': struct_weight, 'tot_num':tot_num} + return {'x': x, 'dx': dx, 'da': da, 'n': n, 'E': E, 'F': F, 'S': S, 'sp_idx': sparse_index, 'struct_weight': struct_weight, 'tot_num': tot_num} @@ -200,14 +200,14 @@ def gen_sparse_index(nlist): idx = 0 for i,item in enumerate(nlist): for jtem in range(item): - res[0,idx] = i - res[1,idx] = idx + res[0, idx] = i + res[1, idx] = idx idx += 1 return res #Load collate from train, valid dataset -def _load_collate(inputs, logfile, scale_factor, pca, train_dataset, valid_dataset, batch_size=1,my_collate=my_collate): +def _load_collate(inputs, logfile, scale_factor, pca, train_dataset, valid_dataset, batch_size=1, my_collate=my_collate): partial_collate = partial( my_collate, atom_types=inputs['atom_types'], diff --git a/simple_nn_v2/models/neural_network.py b/simple_nn_v2/models/neural_network.py index 58ec8e7..f91211c 100644 --- a/simple_nn_v2/models/neural_network.py +++ b/simple_nn_v2/models/neural_network.py @@ -3,11 +3,11 @@ import shutil class FCNDict(torch.nn.Module): - def __init__(self, nets): super(FCNDict, self).__init__() self.nets = torch.nn.ModuleDict(nets) self.keys = self.nets.keys() + def forward(self, x): assert [item for item in self.nets.keys()].sort() == [item for item in x.keys()].sort() res = {} @@ -36,9 +36,9 @@ def write_lammps_potential(self, filename, inputs, scale_factor=None, pca=None): FIL.write('SYM {}\n'.format(len(params))) for ctem in params: - tmp_types = inputs['atom_types'][int(ctem[1])-1] + tmp_types = inputs['atom_types'][int(ctem[1]) - 1] if int(ctem[0]) > 3: - tmp_types += ' {}'.format(inputs['atom_types'][int(ctem[2])-1]) + tmp_types += ' {}'.format(inputs['atom_types'][int(ctem[2]) - 1]) if len(ctem) != 7: raise ValueError("params file must have lines with 7 columns.") @@ -97,7 +97,7 @@ def write_lammps_potential(self, filename, inputs, scale_factor=None, pca=None): FIL.write('LAYER {} {}\n'.format(j+joffset, acti)) - for k in range(nodes[j+joffset]): + for k in range(nodes[j + joffset]): FIL.write('w{} {}\n'.format(k, ' '.join(weights[j][k,:].astype(np.str)))) FIL.write('b{} {}\n'.format(k, biases[j][k])) @@ -106,13 +106,13 @@ def write_lammps_potential(self, filename, inputs, scale_factor=None, pca=None): FIL.close() class FCN(torch.nn.Module): - def __init__(self, dim_input, dim_hidden, acti_func='sigmoid',dropout=None): + def __init__(self, dim_input, dim_hidden, acti_func='sigmoid', dropout=None): super(FCN, self).__init__() self.lin = torch.nn.Sequential() dim_in = dim_input - for i,hn in enumerate(dim_hidden): + for i, hn in enumerate(dim_hidden): if dropout: self.lin.add_module(f'drop_{i}', torch.nn.Dropout(p=dropout)) self.lin.add_module(f'lin_{i}', torch.nn.Linear(dim_in, hn)) @@ -157,7 +157,7 @@ def _read_until(fil, stop_tag): weights = dict() with open(filename) as fil: - atom_types = fil.readline().replace('\n','').split()[1:] + atom_types = fil.readline().replace('\n', '').split()[1:] for item in atom_types: weights[item] = dict() diff --git a/simple_nn_v2/models/run.py b/simple_nn_v2/models/run.py index e5f5c03..8f6cf01 100644 --- a/simple_nn_v2/models/run.py +++ b/simple_nn_v2/models/run.py @@ -5,24 +5,23 @@ import time import numpy as np -from sklearn.decomposition import PCA -from .neural_network import FCNDict, FCN, read_lammps_potential -from .data_handler import StructlistDataset, FilelistDataset, _set_struct_dict, _load_collate, filename_collate -from .train import train, save_checkpoint, _show_structure_rmse, _save_nnp_result, _save_atomic_E +from simple_nn_v2.models.neural_network import FCNDict, FCN, read_lammps_potential +from simple_nn_v2.models.data_handler import StructlistDataset, FilelistDataset, _set_struct_dict, _load_collate, filename_collate +from simple_nn_v2.models.train import train, save_checkpoint, _show_structure_rmse, _save_nnp_result, _save_atomic_E -def train_NN(inputs, logfile,user_optimizer=None): - #Set default type +def train_NN(inputs, logfile, user_optimizer=None): if inputs['neural_network']['double_precision']: torch.set_default_dtype(torch.float64) - # Initialize model - model, optimizer, criterion, scale_factor, pca = _init_model(inputs, logfile, user_optimizer=None) + device = _set_device(logfile) + model = _initialize_model(inputs, logfile, device) + optimizer = _initialize_optimizer(inputs, model, logfile) + criterion = torch.nn.MSELoss(reduction='none').to(device=device) # Resume job if possible and load dataset & scale_factor, pca - scale_factor, pca, train_dataset, valid_dataset, best_loss = _load_data(\ - inputs, logfile, model, optimizer, scale_factor, pca) + scale_factor, pca, train_dataset, valid_dataset, best_loss = _load_data(inputs, logfile, model, optimizer) # Load data loader if inputs['neural_network']['full_batch']: @@ -39,16 +38,14 @@ def train_NN(inputs, logfile,user_optimizer=None): _do_train(inputs, logfile, train_loader, valid_loader, model, optimizer, criterion, scale_factor, pca, best_loss, train_struct_dict, valid_struct_dict) # End of progran & Best loss witten - -#Initialize model with input, and set default value scale, PCA -def _init_model(inputs, logfile, user_optimizer=None): +def _initialize_model(inputs, logfile, device): #Set default configuration - logfile.write('Initialize pytorch model\n') + logfile.write("Initialize pytorch model\n") model = {} for item in inputs['atom_types']: temp_nodes = [int(jtem) for jtem in inputs['neural_network']['nodes'].split('-')] #Extract number of symmetry functions in params - with open(inputs['descriptor']['params'][item],'r') as f: + with open(inputs['descriptor']['params'][item], 'r') as f: tmp_symf = f.readlines() sym_num = len(tmp_symf) #Make NN for each element @@ -61,6 +58,11 @@ def _init_model(inputs, logfile, user_optimizer=None): logfile.write(weight_log) model = FCNDict(model) #Make full model with elementized dictionary model + model.to(device=device) + + return model + +def _initialize_optimizer(inputs, model, user_optimizer=None): regularization = float(inputs['neural_network']['regularization']) if inputs['neural_network']['method'] == 'Adam': @@ -78,50 +80,50 @@ def _init_model(inputs, logfile, user_optimizer=None): else: raise ValueError - device = "cuda" if torch.cuda.is_available() else "cpu" + return optimizer +def _set_device(logfile): + device = 'cuda' if torch.cuda.is_available() else 'cpu' logfile.write(f"Use {device} in model\n") device = torch.device(device) - model.to(device=device) - criterion = torch.nn.MSELoss(reduction='none').to(device=device) + return device + +#Load data (pca,scale fator, dataset) , resume from checkpoint +def _load_data(inputs, logfile, model, optimizer): scale_factor = None pca = None - - return model, optimizer, criterion, scale_factor, pca -#Load data (pca,scale fator, dataset) , resume from checkpoint -def _load_data(inputs, logfile, model, optimizer, scale_factor, pca): - if inputs["neural_network"]["continue"]: - logfile.write('Load pytorch model from {0}\n'.format(inputs["neural_network"]["continue"])) - checkpoint = torch.load(inputs["neural_network"]["continue"]) + if inputs['neural_network']['continue']: + logfile.write("Load pytorch model from {0}\n".format(inputs['neural_network']['continue'])) + checkpoint = torch.load(inputs['neural_network']['continue']) model.load_state_dict(checkpoint['model']) # Load model - if not inputs["neural_network"]['clear_prev_network']: - logfile.write('Load previous model, optimizer, learning rate\n') + if not inputs['neural_network']['clear_prev_network']: + logfile.write("Load previous model, optimizer, learning rate\n") #OPTIMIZER optimizer.load_state_dict(checkpoint['optimizer']) for pg in optimizer.param_groups: pg['lr'] = inputs['neural_network']['learning_rate'] if not inputs['neural_network']['clear_prev_status']: - logfile.write('Load previous loss : {0:6.2e}, epoch : {1}\n'.format(checkpoint['loss'], checkpoint['epoch'])) + logfile.write("Load previous loss : {0:6.2e}, epoch : {1}\n".format(checkpoint['loss'], checkpoint['epoch'])) loss = checkpoint['loss'] inputs['neural_network']['start_epoch'] = checkpoint['epoch'] else: loss = float('inf') # Read LAMMPS potential - if inputs["neural_network"]['read_potential']: + if inputs['neural_network']['read_potential']: _load_lammps_potential(inputs, logfile, model) # Load scale file & PCA file if inputs['descriptor']['calc_scale']: - logfile.write('Load scale from checkpoint\n') + logfile.write("Load scale from checkpoint\n") scale_factor = checkpoint['scale_factor'] if inputs['neural_network']['pca']: - logfile.write('Load pca from checkpoint\n') + logfile.write("Load pca from checkpoint\n") pca = checkpoint['pca'] else: #Not resume, load scale_factor, pca pytorch savefile in preprocess @@ -129,11 +131,11 @@ def _load_data(inputs, logfile, model, optimizer, scale_factor, pca): loss = float('inf') if inputs['descriptor']['calc_scale']: scale_factor = torch.load('./scale_factor') - logfile.write('Scale factor data loaded\n') + logfile.write("Scale factor data loaded\n") if inputs['neural_network']['pca']: pca = torch.load('./pca') - logfile.write('PCA data loaded\n') + logfile.write("PCA data loaded\n") _convert_to_tensor(inputs, logfile, scale_factor, pca) @@ -147,18 +149,18 @@ def _load_data(inputs, logfile, model, optimizer, scale_factor, pca): valid_dataset = FilelistDataset(inputs['descriptor']['valid_list']) try: #Check valid dataset exist valid_dataset[0] - logfile.write('Train & Valid dataset loaded\n') + logfile.write("Train & Valid dataset loaded\n") except: #No validation valid_dataset = None - logfile.write('Train dataset loaded, No valid set loaded\n') + logfile.write("Train dataset loaded, No valid set loaded\n") else: - logfile.write('Test dataset loaded\n') + logfile.write("Test dataset loaded\n") #model, optimizer no need to be returned return scale_factor, pca, train_dataset, valid_dataset, loss #Main traning part -def _do_train(inputs, logfile, train_loader, valid_loader, model, optimizer, criterion, scale_factor, pca, best_loss, train_struct_dict = None, valid_struct_dict = None): +def _do_train(inputs, logfile, train_loader, valid_loader, model, optimizer, criterion, scale_factor, pca, best_loss, train_struct_dict=None, valid_struct_dict=None): #Get start time start_time = time.time() @@ -167,7 +169,7 @@ def _do_train(inputs, logfile, train_loader, valid_loader, model, optimizer, cri max_len = len(train_loader) total_epoch = int(inputs['neural_network']['total_epoch']) total_iter = int(inputs['neural_network']['total_epoch']*max_len) - logfile.write('Total training iteration : {0} , epoch : {1}, batch number : {2}\n'.format(total_iter, total_epoch, max_len)) + logfile.write("Total training iteration : {0} , epoch : {1}, batch number : {2}\n".format(total_iter, total_epoch, max_len)) #Learning rate decay schedular if inputs['neural_network']['lr_decay']: @@ -197,8 +199,8 @@ def _do_train(inputs, logfile, train_loader, valid_loader, model, optimizer, cri #Evaluation model if inputs['neural_network']['test']: - print('Evaluation(Testing) model') - logfile.write('Evaluation(Testing) model \n') + print("Evaluation(Testing) model") + logfile.write("Evaluation(Testing) model \n") loss = train(inputs, logfile, train_loader, model, criterion=criterion, start_time=start_time, test=True) if inputs['neural_network']['print_structure_rmse'] and train_struct_dict: _show_structure_rmse(inputs, logfile, train_struct_dict, valid_struct_dict, model, optimizer=optimizer, criterion=criterion, test=True) @@ -216,7 +218,7 @@ def _do_train(inputs, logfile, train_loader, valid_loader, model, optimizer, cri loss = t_loss #Structure rmse part - if (epoch % inputs['neural_network']['show_interval'] == 0) and\ + if (epoch % inputs['neural_network']['show_interval'] == 0) and\ inputs['neural_network']['print_structure_rmse'] and train_struct_dict: _show_structure_rmse(inputs, logfile, train_struct_dict, valid_struct_dict, model, optimizer=optimizer, criterion=criterion) @@ -229,29 +231,29 @@ def _do_train(inputs, logfile, train_loader, valid_loader, model, optimizer, cri best_loss = loss best_epoch = epoch #Checkpoint in model traning when best model - save_checkpoint(epoch, loss, model, optimizer, pca, scale_factor, filename = 'bestmodel.pth.tar' ) + save_checkpoint(epoch, loss, model, optimizer, pca, scale_factor, filename='bestmodel.pth.tar' ) #Checkpoint for save iteration - if inputs['neural_network']['checkpoint_interval'] and (epoch % inputs['neural_network']['checkpoint_interval'] == 0): - save_checkpoint(epoch, loss, model, optimizer, pca, scale_factor, filename = f'epoch_{epoch}.pth.tar') + if inputs['neural_network']['checkpoint_interval'] and (epoch % inputs['neural_network']['checkpoint_interval'] == 0): + save_checkpoint(epoch, loss, model, optimizer, pca, scale_factor, filename=f'epoch_{epoch}.pth.tar') elif not inputs['neural_network']['checkpoint_interval']: save_checkpoint(epoch, loss, model, optimizer, pca, scale_factor) #Save result if inputs['neural_network']['save_result'] and (epoch % inputs['neural_network']['save_interval'] == 0): res_dict = _save_nnp_result(inputs, model, trainset_saved, validset_saved) - logfile.write(f'DFT, NNP result saved at {epoch}\n') + logfile.write(f"DFT, NNP result saved at {epoch}\n") torch.save(res_dict, 'saved_result') #LAMMPS potential save part breaksignal = _save_lammps(inputs, logfile, model, is_best, epoch, scale_factor, pca, err_dict) if breaksignal: - logfile.write('Break point reached. Terminating traning model\n') - save_checkpoint(epoch, loss, model, optimizer, pca, scale_factor, filename = 'breakpoint.pth.tar') + logfile.write("Break point reached. Terminating traning model\n") + save_checkpoint(epoch, loss, model, optimizer, pca, scale_factor, filename='breakpoint.pth.tar') break #End of traning loop : best loss potential written - logfile.write('Best loss lammps potential written at {0} epoch\n'.format(best_epoch)) + logfile.write("Best loss lammps potential written at {0} epoch\n".format(best_epoch)) if inputs['descriptor']['add_NNP_ref']: _save_atomic_E(inputs, logfile, model, trainset_saved, validset_saved) @@ -261,13 +263,13 @@ def _save_lammps(inputs, logfile, model, is_best, epoch, scale_factor, pca, err_ #Lammps save part #Save best lammps potential if set save_best if is_best: - model.write_lammps_potential(filename ='./potential_saved_best_loss', inputs=inputs, scale_factor=scale_factor, pca=pca) + model.write_lammps_potential(filename='./potential_saved_best_loss', inputs=inputs, scale_factor=scale_factor, pca=pca) #Save lammps potential with save_interval - if epoch % inputs['neural_network']['save_interval'] == 0: - model.write_lammps_potential(filename ='./potential_saved_epoch_{0}'.format(epoch), inputs=inputs, scale_factor=scale_factor, pca=pca) - print('Lammps potential written at {0} epoch\n'.format(epoch)) - logfile.write('Lammps potential written at {0} epoch\n'.format(epoch)) + if epoch % inputs['neural_network']['save_interval'] == 0: + model.write_lammps_potential(filename='./potential_saved_epoch_{0}'.format(epoch), inputs=inputs, scale_factor=scale_factor, pca=pca) + print("Lammps potential written at {0} epoch\n".format(epoch)) + logfile.write("Lammps potential written at {0} epoch\n".format(epoch)) #Break if energy, force, stress is under their criteria @@ -281,11 +283,11 @@ def _save_lammps(inputs, logfile, model, is_best, epoch, scale_factor, pca, err_ if breaksignal: for err_type in err_dict.keys(): - print('Ctirerion met {0} : {1:4.2f} < criteria : {2:4}'.format(err_type,err_dict[err_type][0],err_dict[err_type][1])) - logfile.write('Ctirerion met {0} : {1:4.2f} < criteria : {2:4}\n'.format(err_type,err_dict[err_type][0],err_dict[err_type][1])) - model.write_lammps_potential(filename ='./potential_criterion', inputs=inputs, scale_factor=scale_factor, pca=pca) - print('Criterion lammps potential written'.format(err_type,err_dict[err_type][0],err_dict[err_type][1])) - logfile.write('Criterion lammps potential written\n'.format(err_type,err_dict[err_type][0],err_dict[err_type][1])) + print("Ctirerion met {0} : {1:4.2f} < criteria : {2:4}".format(err_type, err_dict[err_type][0], err_dict[err_type][1])) + logfile.write("Ctirerion met {0} : {1:4.2f} < criteria : {2:4}\n".format(err_type, err_dict[err_type][0], err_dict[err_type][1])) + model.write_lammps_potential(filename='./potential_criterion', inputs=inputs, scale_factor=scale_factor, pca=pca) + print("Criterion lammps potential written".format(err_type, err_dict[err_type][0], err_dict[err_type][1])) + logfile.write("Criterion lammps potential written\n".format(err_type, err_dict[err_type][0], err_dict[err_type][1])) return breaksignal @@ -296,13 +298,13 @@ def _check_criteria(inputs,logfile): #Dictionaly with list [error , criteria] err_dict = dict() err_dict['e_err'] = [float('inf') , float(inputs['neural_network']['energy_criteria'])] - logfile.write('Energy criteria used : {0:4} \n'.format(float(inputs['neural_network']['energy_criteria']))) + logfile.write("Energy criteria used : {0:4} \n".format(float(inputs['neural_network']['energy_criteria']))) if inputs['neural_network']['force_criteria'] and inputs['neural_network']['use_force']: err_dict['f_err'] = [float('inf'), float(inputs['neural_network']['force_criteria'])] - logfile.write('Force criteria used : {0:4}\n'.format(float(inputs['neural_network']['force_criteria']))) + logfile.write("Force criteria used : {0:4}\n".format(float(inputs['neural_network']['force_criteria']))) if inputs['neural_network']['stress_criteria'] and inputs['neural_network']['use_stress']: err_dict['s_err'] = [float('inf'), float(inputs['neural_network']['stress_criteria'])] - logfile.write('Stress criteria used : {0:4}\n'.format(float(inputs['neural_network']['stress_criteria']))) + logfile.write("Stress criteria used : {0:4}\n".format(float(inputs['neural_network']['stress_criteria']))) else: err_dict = None @@ -334,8 +336,8 @@ def _load_structure(inputs, logfile, scale_factor, pca): batch_size = len(train_struct_dict[t_key]) else: batch_size = inputs['neural_network']['batch_size'] - train_struct_dict[t_key], valid_struct_dict[t_key] =_load_collate(inputs, - logfile, scale_factor, pca, train_struct_dict[t_key], valid_struct_dict[t_key],batch_size=batch_size) + train_struct_dict[t_key], valid_struct_dict[t_key] = _load_collate(inputs, + logfile, scale_factor, pca, train_struct_dict[t_key], valid_struct_dict[t_key], batch_size=batch_size) return train_struct_dict, valid_struct_dict @@ -343,7 +345,6 @@ def _load_structure(inputs, logfile, scale_factor, pca): def _init_weight(inputs, model): #Get parameter for initialization try: - init_dic = inputs['neural_network']['weight_initializer'] init_name = init_dic['type'] init_params = init_dic['params'] @@ -378,7 +379,7 @@ def _init_weight(inputs, model): init.xavier_unifrom_(lin.bias, gain=gain) elif init_name == 'xavier normal': - weight_log = "Xavier Uniform weight initialization : gain {0:4.2f}\n".format(gain) + weight_log = "Xavier Normal weight initialization : gain {0:4.2f}\n".format(gain) for lin in model.lin: if lin == Linear: init.xavier_normal_(lin.weight, gain=gain) @@ -432,20 +433,20 @@ def _convert_to_tensor(inputs, logfile, scale_factor, pca): max_plus_min = torch.tensor(scale_factor[item][0,:]) max_minus_min = torch.tensor(scale_factor[item][1,:]) scale_factor[item] = [max_plus_min, max_minus_min] #To list format - logfile.write('Convert {0} scale_factor to tensor\n'.format(item)) + logfile.write("Convert {0} scale_factor to tensor\n".format(item)) if inputs['neural_network']['pca']: pca[item][0] = torch.tensor(pca[item][0]) pca[item][1] = torch.tensor(pca[item][1]) pca[item][2] = torch.tensor(pca[item][2]) - logfile.write('Convert {0} PCA to tensor\n'.format(item)) + logfile.write("Convert {0} PCA to tensor\n".format(item)) def _load_lammps_potential(inputs, logfile, model): - logfile.write('Load parameters from lammps potential filename: {0}\n'.format(inputs["neural_network"]['continue'])) - potential_params = read_lammps_potential(inputs["neural_network"]['continue']) + logfile.write("Load parameters from lammps potential filename: {0}\n".format(inputs['neural_network']['continue'])) + potential_params = read_lammps_potential(inputs['neural_network']['continue']) assert [item for item in model.keys].sort() == [item for item in potential_params.keys()].sort() for item in potential_params.keys(): for name, lin in model.nets[item].lin.named_modules(): if name in potential_params[item].keys(): - lin.weight = Parameter(torch.transpose(torch.tensor(potential_params[item][name]['weight']),-1,0)) - lin.bias = Parameter(torch.transpose(torch.tensor(potential_params[item][name]['bias']),-1,0)) + lin.weight = Parameter(torch.transpose(torch.tensor(potential_params[item][name]['weight']), -1,0)) + lin.bias = Parameter(torch.transpose(torch.tensor(potential_params[item][name]['bias']), -1,0)) diff --git a/simple_nn_v2/models/train.py b/simple_nn_v2/models/train.py index 9a1c0c3..ab86925 100644 --- a/simple_nn_v2/models/train.py +++ b/simple_nn_v2/models/train.py @@ -2,8 +2,8 @@ import shutil import time from ase import units -from ..utils.Logger import AverageMeter, ProgressMeter, TimeMeter, StructureMeter -from .data_handler import FilelistDataset, _load_collate +from simple_nn_v2.utils.Logger import AverageMeter, ProgressMeter, TimeMeter, StructureMeter +from simple_nn_v2.models.data_handler import FilelistDataset, _load_collate #This function train NN def train(inputs, logfile, data_loader, model, optimizer=None, criterion=None, scheduler=None, epoch=0, valid=False, err_dict=None,start_time=None, test=False): @@ -18,7 +18,7 @@ def train(inputs, logfile, data_loader, model, optimizer=None, criterion=None, s max_len = len(data_loader) #Training part - for i,item in enumerate(data_loader): + for i, item in enumerate(data_loader): progress_dict['data_time'].update(time.time() - end) loss = _loop_for_loss(inputs, item, model, criterion, progress_dict, struct_weight=(not valid)) @@ -51,9 +51,9 @@ def _init_meters(model, data_loader, optimizer, epoch, valid, use_force, use_str data_time = AverageMeter('data', ':6.3f') losses = AverageMeter('loss', ':8.4e') e_err = AverageMeter('E err', ':6.4e', sqrt=True) - total_time = TimeMeter('total time',':8.4e') + total_time = TimeMeter('total time', ':8.4e') progress_list = [losses, e_err] - progress_dict = {'batch_time': batch_time, 'data_time': data_time, 'losses': losses, 'e_err': e_err, 'total_time':total_time} + progress_dict = {'batch_time': batch_time, 'data_time': data_time, 'losses': losses, 'e_err': e_err, 'total_time': total_time} if use_force: f_err = AverageMeter('F err', ':6.4e', sqrt=True) @@ -98,14 +98,14 @@ def _init_meters(model, data_loader, optimizer, epoch, valid, use_force, use_str def _show_structure_rmse(inputs, logfile, train_struct_dict, valid_struct_dict, model, optimizer=None, criterion=None, cuda=False, test=False): for t_key in train_struct_dict.keys(): log_train = _struct_log(inputs, train_struct_dict[t_key], model, optimizer=optimizer, criterion=criterion, cuda=cuda, test=test) - log_train = "[{0:8}] ".format(t_key)+log_train + log_train = "[{0:8}] ".format(t_key) + log_train if valid_struct_dict[t_key]: - log_valid = _struct_log(inputs, valid_struct_dict[t_key], model,valid=True, optimizer=optimizer, criterion=criterion, cuda=cuda) - log_valid = "\n[{0:8}] ".format(t_key)+log_valid + log_valid = _struct_log(inputs, valid_struct_dict[t_key], model, valid=True, optimizer=optimizer, criterion=criterion, cuda=cuda) + log_valid = "\n[{0:8}] ".format(t_key) + log_valid else: log_valid = "" - outdict = log_train+log_valid + outdict = log_train + log_valid print(outdict) logfile.write(outdict+'\n') @@ -122,7 +122,7 @@ def _struct_log(inputs, data_loader, model, valid=False, optimizer=None, criteri progress_dict = {'losses': losses, 'e_err': e_err} if inputs['neural_network']['use_force']: - f_err = StructureMeter('F err', ':6.4e',sqrt=True) + f_err = StructureMeter('F err', ':6.4e', sqrt=True) progress_list.append(f_err) progress_dict['f_err'] = f_err if inputs['neural_network']['use_stress']: @@ -200,9 +200,9 @@ def _loop_for_loss(inputs, item, model, criterion, progress_dict, struct_weight= item['sp_idx'][atype].size(1))).to_dense(), axis=1) n_atoms += item['n'][atype].to(device=device, non_blocking=cuda) #Energy loss + structure_weight - e_loss = criterion(E_.squeeze()/n_atoms, item['E'].type(dtype).to(device=device, non_blocking=cuda)/n_atoms) - print_e_loss = torch.mean(e_loss*weight) - e_loss = torch.mean(e_loss*weight) + e_loss = criterion(E_.squeeze() / n_atoms, item['E'].type(dtype).to(device=device, non_blocking=cuda) / n_atoms) + print_e_loss = torch.mean(e_loss * weight) + e_loss = torch.mean(e_loss * weight) #Loop for force, stress if inputs['neural_network']['use_force'] or inputs['neural_network']['use_stress']: @@ -256,7 +256,7 @@ def _loop_for_loss(inputs, item, model, criterion, progress_dict, struct_weight= batch_idx = 0 for n in range(n_batch): #Make structure_weighted force tmp_idx = item['tot_num'][n].item() - f_loss[batch_idx:(batch_idx+tmp_idx)] = f_loss[batch_idx:(batch_idx+tmp_idx)]*weight[n].item() + f_loss[batch_idx:(batch_idx+tmp_idx)] = f_loss[batch_idx:(batch_idx+tmp_idx)] * weight[n].item() batch_idx += tmp_idx f_loss = torch.mean(f_loss) loss += inputs['neural_network']['force_coeff'] * f_loss @@ -269,7 +269,7 @@ def _loop_for_loss(inputs, item, model, criterion, progress_dict, struct_weight= batch_idx = 0 for n in range(n_batch): #Make structure_weighted force tmp_idx = item['tot_num'][n].item() - s_loss[batch_idx:(batch_idx+tmp_idx)] = s_loss[batch_idx:(batch_idx+tmp_idx)]*weight[n].item() + s_loss[batch_idx:(batch_idx+tmp_idx)] = s_loss[batch_idx:(batch_idx+tmp_idx)] * weight[n].item() batch_idx += tmp_idx s_loss = torch.mean(s_loss) loss += inputs['neural_network']['stress_coeff'] * s_loss @@ -298,10 +298,10 @@ def _save_nnp_result(inputs, model, train_loader, valid_loader): #Save NNP energy, force, DFT energy, force to use it model.eval() #Set output - res_dict = {'DFT':dict(),'NNP':dict()} - for datatype in ['train','valid']: - res_dict['DFT'][datatype] = {'E':list(),'F':list()} - res_dict['NNP'][datatype] = {'E':list(),'F':list()} + res_dict = {'DFT': dict(), 'NNP': dict()} + for datatype in ['train', 'valid']: + res_dict['DFT'][datatype] = {'E': list(), 'F': list()} + res_dict['NNP'][datatype] = {'E': list(), 'F': list()} res_dict['tot_num'] = list() _loop_to_save(inputs, model, train_loader, res_dict, save_dict='train') @@ -330,7 +330,7 @@ def _loop_to_save(inputs, model, dataset_loader, res_dict, save_dict='train'): for atype in inputs['atom_types']: x[atype] = item['x'][atype].to(device=device, non_blocking=cuda).requires_grad_(True) if x[atype].size(0) != 0: - atomic_E = torch.sparse.DoubleTensor( + atomic_E = torch.sparse.DoubleTensor( item['sp_idx'][atype].long().to(device=device, non_blocking=cuda), model.nets[atype](x[atype]).squeeze(), size=(item['n'][atype].size(0), item['sp_idx'][atype].size(1))).to_dense() @@ -339,7 +339,7 @@ def _loop_to_save(inputs, model, dataset_loader, res_dict, save_dict='train'): for n in range(n_batch): res_dict['tot_num'].append(item['tot_num'][n].item()) - res_dict['NNP'][save_dict]['E'].append(E_[n].item())) + res_dict['NNP'][save_dict]['E'].append(E_[n].item()) res_dict['DFT'][save_dict]['E'].append(item['E'][n].item()) if inputs['neural_network']['use_force']: @@ -349,7 +349,7 @@ def _loop_to_save(inputs, model, dataset_loader, res_dict, save_dict='train'): dEdG = torch.autograd.grad(torch.sum(E_), x[atype], create_graph=True)[0] tmp_force = list() tmp_idx = 0 - for n,ntem in enumerate(item['n'][atype]): + for n, ntem in enumerate(item['n'][atype]): if ntem != 0: tmp_force.append(torch.einsum('ijkl,ij->kl', item['dx'][atype][n].to(device=device, non_blocking=cuda), dEdG[tmp_idx:(tmp_idx + ntem)])) else: @@ -394,36 +394,36 @@ def _loop_to_save_atomic_E(inputs, model, dataset_loader): n_type[atype] = 0 x[atype] = item['x'][atype].to(device=device, non_blocking=cuda).requires_grad_(True) if x[atype].size(0) != 0: - atomic_E[atype] = torch.sparse.DoubleTensor( + atomic_E[atype] = torch.sparse.DoubleTensor( item['sp_idx'][atype].long().to(device=device, non_blocking=cuda), model.nets[atype](x[atype]).squeeze(), size=(item['n'][atype].size(0), item['sp_idx'][atype].size(1))).to_dense() #Save for f in range(n_batch): - pt_dict = torch.load(item['filename'][f]) + pt_dict = torch.load(item['filename'][f]) del pt_dict['filename'] pt_dict['atomic_E'] = dict() for atype in inputs['atom_types']: if (x[atype].size(0) != 0) and (item['n'][atype][f].item() != 0): - file_atomic_E = atomic_E[atype][f][n_type[atype]:n_type[atype]+item['n'][atype][f]] + file_atomic_E = atomic_E[atype][f][n_type[atype]:n_type[atype] + item['n'][atype][f]] pt_dict['atomic_E'][atype] = file_atomic_E n_type[atype] += item['n'][atype][f] - torch.save(pt_dict,item['filename'][f]) + torch.save(pt_dict, item['filename'][f]) #Not use -def _save_dft_info(inputs, trail_dataset, valid_dataset): +def _save_dft_info(inputs, train_dataset, valid_dataset): dft_dict = { 'total': { - 'energy': list(), - 'force': list(), + 'energy': list(), + 'force': list(), }, 'train': { - 'energy': list(), - 'force': list(), + 'energy': list(), + 'force': list(), }, 'valid': { - 'energy': list(), - 'force': list(), + 'energy': list(), + 'force': list(), } } for _, item in enumerate(train_dataset): diff --git a/simple_nn_v2/utils/Logger.py b/simple_nn_v2/utils/Logger.py index 70ae12a..a870d66 100644 --- a/simple_nn_v2/utils/Logger.py +++ b/simple_nn_v2/utils/Logger.py @@ -23,8 +23,8 @@ def update(self, val, n=1): self.avg = self.sum / self.count if self.sqrt: - self.sqrt_val = self.val**0.5 - self.sqrt_avg = self.avg**0.5 + self.sqrt_val = self.val ** 0.5 + self.sqrt_avg = self.avg ** 0.5 #Show average value : average value @@ -69,7 +69,7 @@ def update(self, val, n=1): self.avg = self.sum / self.count if self.sqrt: - self.sqrt_avg = self.avg**0.5 + self.sqrt_avg = self.avg ** 0.5 #Show sqrt value & average value : batch_value ( average_value ) def __str__(self):