train_loader = spk.AtomsLoader(train, batch_size=64)
val_loader = spk.AtomsLoader(val, batch_size=64)
# statistics
atomrefs = dataset.get_atomrefs(properties)
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)
# get statistics
atomrefs = dataset.get_atomrefs(properties)
per_atom = dict(energy=True, forces=False)
means, stddevs = train_loader.get_statistics(properties,
single_atom_ref=atomrefs,
get_atomwise_statistics=per_atom)
# model build
logging.info("build model")
representation = spk.SchNet(n_interactions=6)
output_modules = [
spk.Atomwise(
property="energy",
derivative="forces",
mean=means["energy"],
stddev=stddevs["energy"],
negative_dr=True,
)
]
model = schnetpack.atomistic.model.AtomisticModel(representation,
output_modules)
# build optimizer
optimizer = Adam(params=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),