def build_fn(atom_features: int = 128, message_steps: int = 8, atomwise: bool = False, output_layers: int = 3, reduce_fn: str = 'sum', mean: Optional[float] = None, std: Optional[float] = None): schnet = spk.representation.SchNet( n_atom_basis=atom_features, n_filters=atom_features, n_gaussians=20, n_interactions=message_steps, cutoff=4., cutoff_network=spk.nn.cutoff.CosineCutoff) if atomwise: output = Atomwise(n_in=atom_features, n_layers=output_layers, aggregation_mode=reduce_fn, mean=mean, stddev=std, property='ip') else: output = Moleculewise(n_in=atom_features, n_layers=output_layers, aggregation_mode=reduce_fn, mean=mean, stddev=std, property='ip') return spk.AtomisticModel(representation=schnet, output_modules=output)
def test_charge_correction(schnet_batch, n_atom_basis): """ Test if charge correction yields the desired total charges. """ model = spk.AtomisticModel( spk.SchNet(n_atom_basis), spk.atomistic.DipoleMoment( n_atom_basis, charge_correction="q", contributions="q" ), ) q = torch.randint(0, 10, (schnet_batch["_positions"].shape[0], 1)) schnet_batch.update(q=q) q_i = model(schnet_batch)["q"] assert torch.allclose(q.float(), q_i.sum(1), atol=1e-6)
def get_model(args, train_loader, mean, stddev, atomref, logging=logging): if logging: logging.info("building model...") representation = get_representation(args, train_loader) output_module = get_output_module( args, representation=representation, mean=mean, stddev=stddev, atomref=atomref, ) model = spk.AtomisticModel(representation, [output_module]) if args.parallel: model = nn.DataParallel(model) if logging: logging.info("The model you built has: %d parameters" % spk.utils.count_params(model)) return model
def get_model(args, train_loader, mean, stddev, atomref, logging=None): if args.mode == "train": if logging: logging.info("building model...") if args.dropout == 0 and args.n_layers == 2: from schnetpack.utils import get_representation, get_output_module representation = get_representation(args, train_loader) output_module = get_output_module( args, representation=representation, mean=mean, stddev=stddev, atomref=atomref, ) else: from schnetpack.utils import get_representation #get_output_module# get_representation representation = get_representation(args, train_loader) # representation = get_rep_with_dropout(args, train_loader) # output_module = get_output_module( output_module = get_outmod_with_dropout( args, representation=representation, mean=mean, stddev=stddev, atomref=atomref, ) model = spk.AtomisticModel(representation, [output_module]) if args.parallel: model = nn.DataParallel(model) if logging: logging.info("The model you built has: %d parameters" % spk.utils.count_params(model)) return model else: raise spk.utils.ScriptError("Invalid mode selected: {}".format( args.mode))
def get_model(args, train_loader, mean, stddev, atomref, logging=None): """ Build a model from selected parameters or load trained model for evaluation. Args: args (argsparse.Namespace): Script arguments train_loader (spk.AtomsLoader): loader for training data mean (torch.Tensor): mean of training data stddev (torch.Tensor): stddev of training data atomref (dict): atomic references logging: logger Returns: spk.AtomisticModel: model for training or evaluation """ if args.mode == "train": if logging: logging.info("building model...") representation = get_representation(args, train_loader) output_module = get_output_module( args, representation=representation, mean=mean, stddev=stddev, atomref=atomref, ) model = spk.AtomisticModel(representation, [output_module]) if args.parallel: model = nn.DataParallel(model) if logging: logging.info( "The model you built has: %d parameters" % spk.utils.count_params(model) ) return model else: raise spk.utils.ScriptError("Invalid mode selected: {}".format(args.mode))
def model(args, omdData, atomrefs, means, stddevs): schnet = spk.representation.SchNet( n_atom_basis=args.features, n_filters=args.features, n_gaussians=50, n_interactions=6, cutoff=5.0, cutoff_network=spk.nn.cutoff.CosineCutoff) output_module = get_output_module( args, representation=schnet, mean=means, stddev=stddevs, atomref=atomrefs, ) # output_Bgap = spk.atomistic.Atomwise(n_in=args.features, atomref=atomrefs[OrganicMaterialsDatabase.BandGap], property=OrganicMaterialsDatabase.BandGap, # mean=means[OrganicMaterialsDatabase.BandGap], stddev=stddevs[OrganicMaterialsDatabase.BandGap]) model = spk.AtomisticModel(representation=schnet, output_modules=output_module) if args.parallel: model = nn.DataParallel(model) return model
means, stddevs = train_loader.get_statistics( properties, per_atom=True, atomrefs=atomrefs ) # model build logging.info("build model") representation = spk.SchNet(n_interactions=6) output_modules = [ spk.Atomwise( property=QM9.U0, mean=means[QM9.U0], stddev=stddevs[QM9.U0], atomref=atomrefs[QM9.U0], ) ] model = spk.AtomisticModel(representation, output_modules) # build optimizer optimizer = Adam(model.parameters(), lr=1e-4) # hooks logging.info("build trainer") metrics = [MeanAbsoluteError(p, p) for p in properties] hooks = [CSVHook(log_path=model_dir, metrics=metrics), ReduceLROnPlateauHook(optimizer)] # trainer loss = mse_loss(properties) trainer = Trainer( model_dir, model=model, hooks=hooks,
# Call ShiftedSigmoid for Activation -> You can use, but you do not have to act = ShiftedSigmoid() # Create Model and Optimizer -> Please Note, that here a modified Fork from the original schnetpack is used # This allows a Noise on the Positions, also the reducing Featurevektor in interaction-layer is part of the modification model = schnetpack.representation.SchNet(use_noise=False, noise_mean=0.0, noise_std=0.1, chargeEmbedding = True, ownFeatures = False, nFeatures = 8, finalFeature = None, max_z=200, n_atom_basis=20, n_filters=[32, 24, 16, 8, 4], n_gaussians=25, normalize_filter=False, coupled_interactions=False, trainable_gaussians=False, n_interactions=5, cutoff=2.5, cutoff_network=schnetpack.nn.cutoff.CosineCutoff) # Modification of schnetpack allows an activation-function on the output of the output-network d = schnetpack.atomistic.Atomwise(n_in=20, aggregation_mode='avg', n_layers=4, mode='postaggregate') # , activation=F.relu)#, output_activation=act.forward) model = schnetpack.AtomisticModel(model, d) optimizer = Adam(model.parameters()) # This function counts the number of trainable parameters of the network # Taken from https://stackoverflow.com/questions/48393608/pytorch-network-parameter-calculation (Wasi Ahmad https://stackoverflow.com/users/5352399/wasi-ahmad) def count_parameters(model): paramCount = 0 for parameterName, parameterValue in model.named_parameters(): if parameterValue.requires_grad: n = np.prod(parameterValue.size()) if parameterValue.dim() <= 1: print(parameterName, ':', n) else: layerSize = list(parameterValue.size()) if (len(layerSize) > 0):
def run(split_path,dataset_path,n_train=None,n_val=None,n_epochs=1000): storage_dir="Info" if not os.path.exists(storage_dir): os.makedirs(storage_dir) if os.path.exists(os.path.join(storage_dir,"checkpoints")): shutil.rmtree(os.path.join(storage_dir,"checkpoints")) if os.path.exists(os.path.join(storage_dir,"log.csv")): os.remove(os.path.join(storage_dir,"log.csv")) if os.path.exists(os.path.join(storage_dir,"best_model")): os.remove(os.path.join(storage_dir,"best_model")) data=MD17(dataset_path) atoms,properties=data.get_properties(0) train,val,test=spk.train_test_split( data=data, split_file=split_path, ) train_loader = spk.AtomsLoader(train, batch_size=100, shuffle=True) val_loader = spk.AtomsLoader(val, batch_size=100) means, stddevs = train_loader.get_statistics( spk.datasets.MD17.energy, divide_by_atoms=True ) with open("out.txt","w+") as file: file.write("IN MD17_train") print('Mean atomization energy / atom: {:12.4f} [kcal/mol]'.format(means[MD17.energy][0])) print('Std. dev. atomization energy / atom: {:12.4f} [kcal/mol]'.format(stddevs[MD17.energy][0])) n_features=64 schnet = spk.representation.SchNet( n_atom_basis=n_features, n_filters=n_features, n_gaussians=25, n_interactions=6, cutoff=5., cutoff_network=spk.nn.cutoff.CosineCutoff ) energy_model = spk.atomistic.Atomwise( n_in=n_features, property=MD17.energy, mean=means[MD17.energy], stddev=stddevs[MD17.energy], derivative=MD17.forces, negative_dr=True ) model = spk.AtomisticModel(representation=schnet, output_modules=energy_model) # tradeoff rho_tradeoff = 0.1 optimizer=Adam(model.parameters(),lr=1e-3) # loss function def loss(batch, result): # compute the mean squared error on the energies diff_energy = batch[MD17.energy]-result[MD17.energy] err_sq_energy = torch.mean(diff_energy ** 2) # compute the mean squared error on the forces diff_forces = batch[MD17.forces]-result[MD17.forces] err_sq_forces = torch.mean(diff_forces ** 2) # build the combined loss function err_sq = rho_tradeoff*err_sq_energy + (1-rho_tradeoff)*err_sq_forces return err_sq # set up metrics metrics = [ spk.metrics.MeanAbsoluteError(MD17.energy), spk.metrics.MeanAbsoluteError(MD17.forces) ] # construct hooks hooks = [ trn.CSVHook(log_path=storage_dir, metrics=metrics), trn.ReduceLROnPlateauHook( optimizer, patience=150, factor=0.8, min_lr=1e-6, stop_after_min=True ) ] trainer = trn.Trainer( model_path=storage_dir, model=model, hooks=hooks, loss_fn=loss, optimizer=optimizer, train_loader=train_loader, validation_loader=val_loader, ) # check if a GPU is available and use a CPU otherwise if torch.cuda.is_available(): device = "cuda" else: device = "cpu" # determine number of epochs and train trainer.train( device=device, n_epochs=n_epochs ) os.rename(os.path.join(storage_dir,"best_model"),os.path.join(storage_dir,"model_new"))
def schnet_train_default(self, train_indices, model_path, old_model_path, schnet_args): import schnetpack as spk import schnetpack.train as trn import torch n_val = schnet_args.get("n_val", 100) # LOADING train, val, test if type(train_indices) == int: n_train = train_indices # Preparing storage storage = os.path.join(self.temp_dir, f"schnet_{n_train}") if not os.path.exists(storage): os.mkdir(storage) split_path = os.path.join(storage, "split.npz") train, val, test = spk.train_test_split(data=self.dataset, num_train=n_train, num_val=n_val, split_file=split_path) else: n_train = len(train_indices) # Preparing storage storage = os.path.join(self.temp_dir, f"schnet_{n_train}") if not os.path.exists(storage): os.mkdir(storage) split_path = os.path.join(storage, "split.npz") all_ind = np.arange(len(self.dataset)) # train train_ind = train_indices all_ind = np.delete(all_ind, train_ind) # val val_ind_ind = np.random.choice(np.arange(len(all_ind)), n_val, replace=False) val_ind = all_ind[val_ind_ind] all_ind = np.delete(all_ind, val_ind_ind) split_dict = { "train_idx": train_ind, "val_idx": val_ind, "test_idx": all_ind, } np.savez_compressed(split_path, **split_dict) train, val, test = spk.train_test_split(data=self.dataset, split_file=split_path) print_ongoing_process(f"Preparing SchNet training, {len(train)} points", True) data = self.dataset batch_size = schnet_args.get("batch_size", 10) n_features = schnet_args.get("n_features", 64) n_gaussians = schnet_args.get("n_gaussians", 25) n_interactions = schnet_args.get("n_interactions", 6) cutoff = schnet_args.get("cutoff", 5.0) learning_rate = schnet_args.get("learning_rate", 1e-3) rho_tradeoff = schnet_args.get("rho_tradeoff", 0.1) patience = schnet_args.get("patience", 5) n_epochs = schnet_args.get("n_epochs", 100) # PRINTING INFO i = {} i["batch_size"], i["n_features"] = batch_size, n_features i["n_gaussians"], i["n_interactions"] = n_gaussians, n_interactions i["cutoff"], i["learning_rate"] = cutoff, learning_rate i["rho_tradeoff"], i["patience"] = rho_tradeoff, patience i["n_epochs"], i["n_val"] = n_epochs, n_val print_table("Parameters", None, None, i, width=20) print() train_loader = spk.AtomsLoader(train, shuffle=True, batch_size=batch_size) val_loader = spk.AtomsLoader(val, batch_size=batch_size) # STATISTICS + PRINTS means, stddevs = train_loader.get_statistics("energy", divide_by_atoms=True) print_info( "Mean atomization energy / atom: {:12.4f} [kcal/mol]".format( means["energy"][0])) print_info( "Std. dev. atomization energy / atom: {:12.4f} [kcal/mol]".format( stddevs["energy"][0])) # LOADING MODEL print_ongoing_process("Loading representation and model") schnet = spk.representation.SchNet( n_atom_basis=n_features, n_filters=n_features, n_gaussians=n_gaussians, n_interactions=n_interactions, cutoff=cutoff, cutoff_network=spk.nn.cutoff.CosineCutoff, ) energy_model = spk.atomistic.Atomwise( n_in=n_features, property="energy", mean=means["energy"], stddev=stddevs["energy"], derivative="forces", negative_dr=True, ) model = spk.AtomisticModel(representation=schnet, output_modules=energy_model) print_ongoing_process("Loading representation and model", True) # OPTIMIZER AND LOSS print_ongoing_process("Defining loss function and optimizer") from torch.optim import Adam optimizer = Adam(model.parameters(), lr=learning_rate) def loss(batch, result): # compute the mean squared error on the energies diff_energy = batch["energy"] - result["energy"] err_sq_energy = torch.mean(diff_energy**2) # compute the mean squared error on the forces diff_forces = batch["forces"] - result["forces"] err_sq_forces = torch.mean(diff_forces**2) # build the combined loss function err_sq = rho_tradeoff * err_sq_energy + (1 - rho_tradeoff) * err_sq_forces return err_sq print_ongoing_process("Defining loss function and optimizer", True) # METRICS AND HOOKS print_ongoing_process("Setting up metrics and hooks") metrics = [ spk.metrics.MeanAbsoluteError("energy"), spk.metrics.MeanAbsoluteError("forces"), ] hooks = [ trn.CSVHook(log_path=storage, metrics=metrics), trn.ReduceLROnPlateauHook(optimizer, patience=5, factor=0.8, min_lr=1e-6, stop_after_min=True), ] print_ongoing_process("Setting up metrics and hooks", True) print_ongoing_process("Setting up trainer") trainer = trn.Trainer( model_path=storage, model=model, hooks=hooks, loss_fn=loss, optimizer=optimizer, train_loader=train_loader, validation_loader=val_loader, ) print_ongoing_process("Setting up trainer", True) if torch.cuda.is_available(): device = "cuda" print_info(f"Cuda cores found, training on GPU") else: device = "cpu" print_info(f"No cuda cores found, training on CPU") print_ongoing_process(f"Training {n_epochs} ecpochs, out in {storage}") trainer.train(device=device, n_epochs=n_epochs) print_ongoing_process(f"Training {n_epochs} epochs, out in {storage}", True) os.mkdir(model_path) os.rename(os.path.join(storage, "best_model"), os.path.join(model_path, "model")) shutil.copy(split_path, os.path.join(model_path, "split.npz"))
def create_model(args, atomrefs, means, stddevs, properties, avg_n_atoms): ssp = rescaled_act.ShiftedSoftplus(beta=args.beta) kernel_conv = create_kernel_conv( cutoff=args.rad_maxr, n_bases=args.rad_nb, n_neurons=args.rad_h, n_layers=args.rad_L, act=ssp, radial_model=args.radial_model ) sp = rescaled_act.Softplus(beta=args.beta) if args.res: net = ResNetwork( kernel_conv=kernel_conv, embed=args.embed, l0=args.l0, l1=args.l1, l2=args.l2, l3=args.l3, L=args.L, scalar_act=sp, gate_act=rescaled_act.sigmoid, avg_n_atoms=avg_n_atoms ) else: net = Network( kernel_conv=kernel_conv, embed=args.embed, l0=args.l0, l1=args.l1, l2=args.l2, l3=args.l3, L=args.L, scalar_act=sp, gate_act=rescaled_act.sigmoid, avg_n_atoms=avg_n_atoms ) ident = torch.nn.Identity() if args.mlp_out: outnet = OutputMLPNetwork( kernel_conv=kernel_conv, previous_Rs=net.Rs[-1], l0=args.outnet_l0, l1=args.outnet_l1, l2=args.outnet_l2, l3=args.outnet_l3, L=args.outnet_L, scalar_act=sp, gate_act=rescaled_act.sigmoid, mlp_h=args.outnet_neurons, mlp_L=args.outnet_layers, avg_n_atoms=avg_n_atoms ) else: outnet = OutputScalarNetwork( kernel_conv=kernel_conv, previous_Rs=net.Rs[-1], scalar_act=ident, avg_n_atoms=avg_n_atoms ) output_modules = [ spk.atomistic.Atomwise( property=prop, mean=means[prop], stddev=stddevs[prop], atomref=atomrefs[prop], outnet=outnet, # aggregation_mode='sum' ) for prop in properties ] model = spk.AtomisticModel(net, output_modules) return model
def train(self, resultfolder, traindb='Data/dataset_10_12_train_combined.db', benchdb='Data/dataset_10_12_test.db', traindata='../../Data/combined1618/', benchdata='../../Data/test/', indexpath='../../Data/INDEX_refined_data.2016.2018', properties=['KD'], threshold=10, cutoff=8, numVal=150, featureset=False, trainBatchsize=8, valBatchsize=1, benchBatchsize=1, natoms=None, props=False, ntrain=4444, ntest=290, use_noise=False, noise_mean=0.0, noise_std=0.1, chargeEmbedding=True, ownFeatures=False, nFeatures=8, finalFeature=None, max_z=200, n_atom_basis=20, n_filters=32, n_gaussians=25, normalize_filter=False, coupled_interactions=False, trainable_gaussians=False, n_interactions=5, distanceCutoff=2.5, cutoff_network=schnetpack.nn.cutoff.CosineCutoff, outputIn=32, outAggregation='avg', outLayer=2, outMode='postaggregate', outAct=schnetpack.nn.activations.shifted_softplus, outOutAct=None, n_acc_steps=8, remember=10, ensembleModel=False, n_epochs=150, lr=1e-3, weight_decay=0, train_loader=None, val_loader=None, splitfile=None, noProtons=False): print('Device: ', torch.cuda.current_device()) torch.cuda.empty_cache() # Define Folder for Results Resultfolder = resultfolder #You can activate shifted sigmoid #act = ShiftedSigmoid() if train_loader is None or val_loader is None: f = open("log.txt", "a") f.writelines( str(datetime.datetime.now()) + ' ' + Resultfolder + ' create loader by its own' + '\n') f.close() train_loader, val_loader, bench_loader = self.createDataloader( traindb=traindb, benchdb=benchdb, traindata=traindata, benchdata=benchdata, indexpath=indexpath, properties=properties, threshold=threshold, cutoff=cutoff, numVal=numVal, featureset=featureset, trainBatchsize=trainBatchsize, valBatchsize=valBatchsize, benchBatchsize=benchBatchsize, natoms=natoms, props=props, ntrain=ntrain, ntest=ntest, splitfile=splitfile, noProtons=noProtons) model = schnetpack.representation.SchNet( use_noise=use_noise, noise_mean=noise_mean, noise_std=noise_std, chargeEmbedding=chargeEmbedding, ownFeatures=ownFeatures, nFeatures=nFeatures, finalFeature=finalFeature, max_z=max_z, n_atom_basis=n_atom_basis, n_filters=n_filters, n_gaussians=n_gaussians, normalize_filter=normalize_filter, coupled_interactions=coupled_interactions, trainable_gaussians=trainable_gaussians, n_interactions=n_interactions, cutoff=distanceCutoff, cutoff_network=cutoff_network) d = schnetpack.atomistic.Atomwise(n_in=outputIn, aggregation_mode=outAggregation, n_layers=outLayer, mode=outMode, activation=outAct, output_activation=outOutAct) model = schnetpack.AtomisticModel(model, d) optimizer = AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) print(model) print('number of trainable parameters =', SchnetTraining.count_parameters(model)) # Defines Metrics, Hooks and Loss loss = SchnetTraining.mse_loss() # MSE Metrics for Validation metrics = [schnetpack.metrics.MeanSquaredError('KD', model_output='y')] # CSVHook -> Creates a log-file # ReduceLROnPlateauHook -> LR-Decay hooks = [ schnetpack.train.CSVHook(log_path=Resultfolder, metrics=metrics), schnetpack.train.ReduceLROnPlateauHook(optimizer, patience=7, factor=0.5, min_lr=1e-5), schnetpack.train.EarlyStoppingHook(patience=40, threshold_ratio=0.001) ] # Create Trainer -> n_acc_steps is for accumulating gradient, if the graphiccard can not handle big batch-sizes trainer = schnetpack.train.Trainer(model_path=Resultfolder, model=model, loss_fn=loss, train_loader=train_loader, optimizer=optimizer, validation_loader=val_loader, hooks=hooks, n_acc_steps=n_acc_steps, keep_n_checkpoints=10, checkpoint_interval=5, remember=remember, ensembleModel=ensembleModel) # Start training on cuda for infinite epochs -> Use early stopping trainer.train(device='cuda', n_epochs=n_epochs)
def atomistic_model(schnet, output_modules): return spk.AtomisticModel(schnet, output_modules)
cutoff=SchNet_cutoff, cutoff_network=spk.nn.cutoff.CosineCutoff, trainable_gaussians=trainable_gaussians, ) # %% energy_model = spk.atomistic.Atomwise(n_in=n_features, property='energy', mean=means['energy'], stddev=stddevs['energy'], derivative='forces', negative_dr=True) # %% model = spk.AtomisticModel(representation=schnet, output_modules=energy_model) # %% Multi-GPUs import torch if torch.cuda.device_count() > 1: print("Let's use", torch.cuda.device_count(), "GPUs!") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs model = torch.nn.DataParallel(model) # %% import torch # loss function the same as https://aip.scitation.org/doi/pdf/10.1063/1.5019779 def loss(batch, result): # compute the mean squared error on the energies