SchNetPack: training atomistic neural networks through a Hydra CLI

The problem SchNetPack solves is configuration, not model code

Writing a message-passing neural network that predicts a potential energy surface is not the hard part anymore. The hard part is the surrounding scaffolding: deciding which properties the model outputs, how each output is weighted in the loss, which metrics are tracked, how the dataset is split and batched, and how a run is reproduced six months later. SchNetPack's answer is to make all of that declarative. The README describes it as a toolbox that contains basic building blocks of atomistic neural networks, manages their training, and provides simple access to common benchmark datasets. The audience is therefore a researcher or engineer who already knows what a continuous-filter convolution is and does not want to rebuild the training loop around it. If you are looking for a library that hands you a pretrained model you can call on an arbitrary XYZ file, this is not that. The shipped entry point is a training script, and the unit of work is an experiment config.

What is actually inside: representations, output modules, and an MD layer

The feature list is more layered than the name suggests. Two representations are listed: SchNet, described as an end-to-end continuous-filter CNN, and PaiNN, described as equivariant message-passing. On top of those sit output modules for dipole moments, polarizability, stress, and general response properties. Below them sit physics modules: electrostatics, Ewald summation, and ZBL repulsion. Separately, there is a GPU-accelerated molecular dynamics code that the README says includes path-integral MD, thermostats and barostats. That last item is the one people overlook. Most atomistic ML packages stop at prediction, and you export forces to an external MD engine. SchNetPack keeps the integrator in the same repository, which matters if your workflow is to train a surrogate and immediately run trajectories on it. The trade-off is a larger dependency surface and a codebase that spans two fairly different concerns: supervised property regression and time integration.

The data flow runs through Hydra config groups into a Lightning module

The architecture visible in the README is a composition of config groups resolved by Hydra and instantiated by PyTorch Lightning. The central class is AtomisticModel, described as a pytorch_lightning.LightningModule. Representations and output modules are selected by config, and the README states that the pre-defined configuration for MD17 selects the representation and output modules that AtomisticModel uses. This is why the same class handles both single-property prediction and potential energy surfaces. The mechanism is the _target_ key. A representation block names a Python class such as schnetpack.representation.PaiNN with its own hyperparameters, and nested radial basis and cutoff function blocks are instantiated the same way. Forces and stress tensors are not separate models; the README explains they come from derivatives of the energy output, which is why the MD17 example works by configuring outputs rather than writing a new class. If you have used Hydra before, this will read as a normal instantiate-style config tree. If you have not, the indirection is the main learning cost.

Getting a run started: spktrain, experiment names, and the dot versus slash rule

Installation is either pip install schnetpack or a git clone followed by pip install . from the repository root. Installing adds the spktrain script to your PATH. The recommended first step is to create a working directory, and the README's example is mkdir spk_workdir followed by cd spk_workdir. From there, spktrain experiment=qm9_atomwise starts a SchNet training run with default settings, downloading QM9 into spk_workdir/data if it is not already present. Runs land under spk_workdir/runs in a subdirectory named by a unique run id hash, and you can redirect them with run.data_dir, run.path and run.id overrides on the same command line. The syntax rule that trips people up is stated plainly in the README: a slash loads a preconfigured config group, a dot changes an individual value. So model/representation=painn swaps the representation wholesale, while model.representation.n_interactions=5 adjusts one field inside whatever representation is already selected. Combining them on one line, as the PaiNN example does, is the normal pattern. Running spktrain experiment=qm9_atomwise --help prints the full config with every parameter that can be changed, which is the fastest way to learn the available keys.

Loss weighting for energy and forces is manual and easy to get wrong

The MD17 example exposes the least forgiving part of the workflow. Because reference energies and forces are both available, both need loss weights, and they are defined inside the task config group as a list of ModelOutput entries. Each entry names a property via globals.energy_key or globals.forces_key, attaches a torch.nn.MSELoss, and registers torchmetrics MeanAbsoluteError and MeanSquaredError for tracking. The README's example config uses a loss_weight of 0.005 for energy and 0.995 for forces, and the accompanying text recommends a stronger weight on forces. Nothing in the framework validates that these weights sum sensibly or that the units are consistent. Get the ratio wrong and training will converge to something that looks fine on one metric and is useless for dynamics. This is a genuine limitation rather than a documentation gap: the config system makes the weights explicit and overridable, which is the right design, but it does not protect you from a bad choice. Expect to run several short trainings to calibrate the ratio before committing GPU time.

Where SchNetPack is the wrong tool

The clearest failure mode is the one the README implies rather than states. Every example is built around a bundled benchmark: QM9, MD17, and the other benchmark datasets the text refers to without enumerating. The value proposition is that a config group already exists for your problem. If your data is a proprietary set of trajectories in a custom format, you are writing a dataset class and a config group from scratch, and at that point you are using SchNetPack mostly as a source of representation modules. That is a legitimate use, but it is a much smaller commitment than adopting the CLI, and it changes what you should evaluate. A second boundary is model scope. The README lists SchNet and PaiNN. If your research depends on a different equivariant architecture, the config groups will not help you, and the Lightning and Hydra layers become overhead you are maintaining rather than infrastructure you are reusing. A third is the licence. The repository metadata reports NOASSERTION, which means the licence could not be automatically classified. That is not a permissive signal, and it is not a restrictive one either. It is an unknown, and unknown licensing is a real blocker for commercial adoption.

The alternative is the template SchNetPack is built on

The README is unusually candid about its own lineage: the CLI is based on the PyTorch Lightning and Hydra template at github.com/ashleve/lightning-hydra-template, and the badge at the top of the README names that template directly. That template is the honest alternative, and the difference is concrete. The template gives you the Hydra config groups, the Lightning training loop, the logger wiring and the run directory convention, with no atomistic content at all. SchNetPack adds the representations, the output modules for dipoles, polarizability and stress, the electrostatics and Ewald and ZBL modules, the benchmark dataset configs, and the GPU molecular dynamics code. So the decision is whether you want the physics and chemistry layer or only the training scaffold. If you have your own representation and your own dataset, the template is smaller and you avoid inheriting a config tree you will mostly override. If you want SchNet or PaiNN on QM9 or MD17 today, the template gives you nothing and SchNetPack gives you a working command.

Maintenance cost and the release cadence

Version numbers and dates are the only maintenance signal available here. The releases listed are v2.1.0 in August 2024, v2.1.1 in September 2024, and v2.2.0 in December 2025, with the last push to master in August 2026. That gap between v2.1.1 and v2.2.0 is roughly fifteen months, so minor releases are not frequent, and pinning a version for a long-running project is reasonable. The upgrade cost is concentrated in the config tree rather than in Python APIs, because Hydra config groups and _target_ paths are the interface. A representation class rename or a moved module path breaks a YAML file, not a function call, and the failure appears at instantiation time when spktrain resolves the config. The practical consequence is that you should keep your experiment configs in version control alongside the SchNetPack version you pin, and treat a version bump as a config migration rather than a dependency update. On licensing, the NOASSERTION classification means you should open the repository's licence file yourself and read it before you build anything you intend to distribute. That is a factual gap in the metadata, not legal advice.

Editorial conclusion

Adopt SchNetPack if you want SchNet or PaiNN on a benchmark dataset and you are willing to learn Hydra's dot-versus-slash override syntax, because that syntax is the actual interface to the library. Do not adopt it if your model is not one of the shipped representations or your dataset is not one of the bundled benchmarks, since the framework's value is concentrated in its preconfigured experiment files. Before committing, run spktrain experiment=qm9_atomwise --help and read the full resolved config, then check the repository's licence file directly, because the metadata reports NOASSERTION and that is not a licence grant.

Official sources

1. atomistic-machine-learning/schnetpack on GitHub
2. Issues
3. README
4. Releases