def build_model(mean, stddev, model_properties, atomrefs, additional_outputs, n_atom_basis, name, cutoff): """ Build the model from a given config. Args: mean (dict): means of the dataset properties stddev (dict): stds of the dataset properties model_properties (list): properties for the model atomrefs: atomic reference data additional_outputs (list): additional model output that is not back-propagated n_atom_basis (int): number of features used to describe atomic environments name (str): choose model representation cutoff (float): cutoff radius of filters Returns: Model object """ if name == 'schnet': representation = build_schnet(return_intermediate=False) else: raise ModelError('Unknown model: {:s}'.format(name)) cutoff_function = get_cutoff() output = PropertyModel(n_atom_basis, model_properties + additional_outputs, mean, stddev, atomrefs, cutoff_network=cutoff_function, cutoff=cutoff) model = AtomisticModel(representation, output) return model
def load_model(modelpath, cuda=True): """ Load a trained model from its directory and prepare it for simulations with ASE. Args: modelpath (str): Path to model directory. cuda (bool): Use cuda (default=True). Returns: object: Model class specified in molecular_dynamics. Contains the model, model type and device. """ # Load stored arguments argspath = os.path.join(modelpath, 'args.json') args = read_from_json(argspath) # Reconstruct model based on arguments if args.model == 'schnet': representation = SchNet(args.features, args.features, args.interactions, args.cutoff, args.num_gaussians) atomwise_output = Energy(args.features, return_force=True, create_graph=True) elif args.model == 'wacsf': # Build HDNN model mode = ('weighted', 'Behler')[args.behler] # Convert element strings to atomic charges elements = frozenset((atomic_numbers[i] for i in sorted(args.elements))) representation = BehlerSFBlock(args.radial, args.angular, zetas=set(args.zetas), cutoff_radius=args.cutoff, centered=args.centered, crossterms=args.crossterms, mode=mode, elements=elements) representation = StandardizeSF(representation, cuda=args.cuda) atomwise_output = ElementalEnergy(representation.n_symfuncs, n_hidden=args.n_nodes, n_layers=args.n_layers, return_force=True, create_graph=True, elements=elements) else: raise ValueError('Unknown model class:', args.model) model = AtomisticModel(representation, atomwise_output) # Load old parameters model.load_state_dict(torch.load(os.path.join(modelpath, 'best_model'))) # Set cuda if requested device = torch.device("cuda" if cuda else "cpu") model = model.to(device) # Store into model wrapper for calculator ml_model = Model(model, args.model, device) return ml_model
def get_model(representation, output_modules, parallelize=False): model = AtomisticModel(representation, output_modules) if parallelize: model = nn.DataParallel(model) logging.info("The model you built has: %d parameters" % compute_params(model)) return model
def get_model(args, train_loader, mean, stddev, atomref, aggregation_mode, 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, aggregation_mode=aggregation_mode, ) model = AtomisticModel(representation, [output_module]) if args.parallel: model = nn.DataParallel(model) if logging: logging.info("The model you built has: %d parameters" % schnetpack.utils.spk_utils.count_params(model)) return model else: raise spk.utils.ScriptError("Invalid mode selected: {}".format( args.mode))
def get_model(train_args, basisdef, orbital_energies, mean, stddev, parallelize): if train_args.cutoff_function == "hard": cutoff_network = HardCutoff elif train_args.cutoff_function == "cosine": cutoff_network = CosineCutoff elif train_args.cutoff_function == "mollifier": cutoff_network = MollifierCutoff num_gaussians = int(train_args.cutoff * 5) dirs = train_args.directions if train_args.directions > 0 else None schnorb = stl.model.SchNOrb(n_factors=train_args.factors, lmax=train_args.lmax, n_interactions=train_args.interactions, directions=dirs, n_cosine_basis=train_args.orbbasis, n_gaussians=num_gaussians, cutoff_network=cutoff_network, coupled_interactions=args.coupled) print('SchNorb params: %.2fM' % (sum(p.numel() for p in schnorb.parameters()) / 1000000.0)) hamiltonian = stl.model.Hamiltonian(basisdef, lmax=train_args.lmax, n_cosine_basis=train_args.orbbasis, directions=dirs, orbital_energies=orbital_energies, quambo=train_args.quambo, mean=mean, stddev=stddev, return_forces=train_args.forces, create_graph=train_args.forces) print('Outnet params: %.2fM' % (sum(p.numel() for p in hamiltonian.parameters()) / 1000000.0)) schnorb = AtomisticModel(schnorb, hamiltonian) if parallelize: schnorb = nn.DataParallel(schnorb) return schnorb
def atomistic_model(schnet, energy_module, property_output): return AtomisticModel(schnet, output_modules=[energy_module, property_output])
def train(self, n_epochs, lr, loss_fn, batch_size, num_workers, device, patience=100, threshold_ratio=0.0001): self.i += 1 reduced = self.dataset.create_subset(self.idx_red) num_val = round(0.10 * len(reduced)) train, val, test = train_test_split(data=reduced, num_train=len(reduced) - num_val, num_val=num_val) train_loader = AtomsLoader(train, batch_size=round(batch_size), num_workers=num_workers, shuffle=True, pin_memory=True) val_loader = AtomsLoader(val, batch_size=round(batch_size / 2), num_workers=num_workers, pin_memory=True) representation = SchNet(n_atom_basis=self.n_atom_basis, n_filters=self.n_filters, n_interactions=self.n_interactions, cutoff=self.cutoff, n_gaussians=self.n_gaussians) output_modules = Atomwise(representation.n_atom_basis, n_layers=self.n_layers, property='energy', derivative='forces', stress='stress', negative_dr=True, create_graph=True) model = AtomisticModel(representation, output_modules) optimizer = Adam(model.parameters(), lr=lr) hooks = [ CSVHook('log_%i' % self.i, [ MeanAbsoluteError('energy', 'energy'), MeanAbsoluteError('forces', 'forces', element_wise=True), MeanAbsoluteError('stress', 'stress'), R2Score('energy', 'energy'), R2Score('forces', 'forces', element_wise=True), R2Score('stress', 'stress') ], every_n_epochs=1) ] hooks.append(EarlyStoppingHook(patience, threshold_ratio)) trainer = Trainer('output_%i/' % self.i, model, loss_fn, optimizer, train_loader, val_loader, hooks=hooks, keep_n_checkpoints=1, checkpoint_interval=n_epochs) print('Running training!') print(' Reduced images: %i' % len(reduced)) print(' Traning images: %i' % len(train)) print(' Validation images: %i' % len(val)) print('') trainer.train(device, n_epochs)