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)