Repository navigation
Final schnetpack 2.x release - #815
Merged
Merged
Conversation
Add `load_task_from_checkpoint`, which passes `weights_only=False` to `load_from_checkpoint` only when the installed PyTorch Lightning supports it, and use it in `spktrain` to load and test the best checkpoint instead of `trainer.test(ckpt_path="best")`. Partial cherry-pick of the checkpoint-loading changes in cli.py and utils/compatibility.py from 11f6f71 (#781); the rest of that PR (the dataset/datamodule refactor) changes the public API and is left on dev. (partially cherry picked from commit 11f6f71) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
* feat: add progress bar callback to training configuration * fix: preserve raw weights when resuming EMA training from a checkpoint Checkpoints save the model right after the EMA average is copied in during validation, so on_fit_start was calling store() on the averaged weights and treating them as raw, causing resumed training to continue from the smoothed average instead of the true raw trajectory. Track whether the average is currently active so resume can restore the actual raw weights first. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * feat: enhance checkpoint handling in training process * black * black old * Revert "fix: preserve raw weights when resuming EMA training from a checkpoint" This reverts commit 00263dd. * chore: pin black version to 24.4.2 in pyproject.toml for consistent formatting --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> (cherry picked from commit 175e9cc)
* feat: configure wandb logger with project/run name Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> (cherry picked from commit 847f666)
Hold the TorchScript module and its device behind an opaque pointer, so that
libtorch's headers are no longer visible to the rest of LAMMPS. Without this
the build fails:
ATen/core/Generator.h:121:3: error: reference to 'Device' is ambiguous
note: candidates are: 'LAMMPS_NS::ExecutionSpace LAMMPS_NS::Device'
note: 'struct c10::Device'
LAMMPS collects every pair-style header into the generated style_pair.h, which
force.cpp and lammps.cpp include *after* their own "using namespace LAMMPS_NS;".
libtorch's headers are then parsed with LAMMPS_NS injected into the global
scope, so the unqualified "Device" that ATen uses matches both c10::Device and
the LAMMPS_NS::Device enumerator of "enum ExecutionSpace" in src/pointers.h.
Only force.cpp and lammps.cpp fail. pair_schnetpack.cpp itself compiles fine,
because it includes its own header before the using-directive -- the bug is
purely one of include order.
Note that the enumerator is not guarded by #ifdef LMP_KOKKOS despite the
comment above it, so disabling KOKKOS does not avoid the clash. The same
collision affects other libtorch-based pair styles, e.g.
mir-group/pair_nequip_allegro#66
The pair style still requires and links libtorch; only header visibility
changes. Behaviour, input syntax and results are unaffected.
Seen with LAMMPS 22Jul2025_update5, PyTorch 2.13.0, GCC 15.2.0, C++17.
(cherry picked from commit 6a4aac7)
(cherry picked from commit 7d9bcf5)
(cherry picked from commit a19928a) Backport note: only the info-level split load/save log messages are taken. The `split_file` change in configs/data/custom.yaml fixes a path introduced by the #781 datamodule refactor and does not apply to master, where the default split file already resolves into the run directory. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
(cherry picked from commit 76e2366)
* fixed multi output bug (cherry picked from commit 8af955f)
…802) PyTorch Lightning's EarlyStopping (check_finite=True) stops training silently when the monitored metric becomes NaN or inf. spktrain then continued with testing and model export, and the run was reported as finished. Now spktrain raises a RuntimeError after fit in that case. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 299fffa)
* new ci workflow * updated pre-commit * new manifest * ruff check passes * ruff * remove build workflow for now * Update activation parameter type hint to allow None Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * ignore docs for now in CI --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> (cherry picked from commit b0806ed) Backport note: on dev this commit sits on top of the #781 datamodule refactor, which is not part of v2_dev. Files that conflict with that refactor (data/, datasets/, cli.py, a few configs, tests and notebooks) were resolved by keeping the v2 code and running this commit's own pre-commit config (ruff v0.14.0) over it, plus the manual lint fixes ruff cannot apply: explicit imports instead of `from schnetpack.data import *` in datasets/, `except Exception` instead of bare `except`, `raise ... from`, and %-format log calls. data/provider.py (added by #781) is dropped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
- pytorch_lightning >=2.6.1,<3 (was >=2.2,<3): before 2.6.1 the DataLoader worker seeds differed between epochs only in their high bits, which torch's CPU generator ignores, so with seed_everything(workers=True), as used by spktrain, random transforms in workers replayed the same random stream every epoch. - torch <3 and numpy <3: cap the major versions, since this release will not be updated for future API breaks. - Remove pytest from the runtime dependencies; it is in the test extra. - Cap pytest <9.2 in the test extra: pytest 9.2 breaks the pytest-benchmark fixture (up to 5.3.0), which reads node._nodeid. - Align the matching bounds in docs/sphinx-requirements.txt. Tested (CPU, test suite plus a spktrain train/resume run) at the lower bounds (torch 2.5.0, pytorch_lightning 2.6.1, numpy 2.0.0, hydra-core 1.3.0, ase 3.23.0) and with the latest releases (torch 2.15.0, pytorch_lightning 2.6.6, numpy 2.5.4, hydra-core 1.3.8, ase 3.29.0). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Since PyTorch Lightning 2.6.0, ModelCheckpoint(save_last=True) only writes the last checkpoint in steps that also save a top-k checkpoint (Lightning-AI/pytorch-lightning#22043). Once the monitored metric stops improving, last.ckpt stays at the best epoch, so resuming spktrain from it silently repeats every epoch since then. schnetpack's ModelCheckpoint now writes the last checkpoint at the end of each validation run (or training epoch with save_on_train_epoch_end) whenever Lightning did not, matching Lightning < 2.6. It never writes twice in one step, so it is a no-op on Lightning versions without the regression and once it is fixed upstream. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
log_hyperparameters sent the config to `trainer.logger`, which is only the first logger, and disabled only that one. Any further logger got the task's hyperparameters from Lightning instead, which hold objects that cannot be written to YAML, so spktrain crashed with two loggers, in either order, e.g. `logger=[tensorboard,csv]`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
PyTorch Lightning 2.0 removed the `gpus` argument of the Trainer, so `trainer=debug_trainer` crashed spktrain at startup. The default trainer it inherits from already sets `devices: 1`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`metrics` defaults to None, but the constructor called `metrics.items()` on it, so creating a ModelOutput without metrics raised an AttributeError. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
config.yaml was written before spktrain drew a random seed, so it held no seed at all and the run could not be repeated from its config. A resumed run also drew a new seed. The config is now written after the seed is chosen, and a resumed run reuses the stored seed unless a new one is given. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`seed` was only read when passed as `+seed=...`, although the docs list it as a top-level key of the train config. It is now `seed: null` in train.yaml, so `seed=42` works as an override; null still draws a random seed, as spkmd's config does. `trainer.log_every_n_steps` is added to the default trainer with Lightning's default of 50. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
TensorBoardLogger opened a new version_N on every spktrain start, so each resume showed up as a separate run. A run directory holds exactly one run, since spktrain resumes whenever it finds a config.yaml there, so the version is now fixed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Lightning does not store RNG states in checkpoints, and spktrain seeds again on restart. A resumed run therefore shuffled the training data differently and diverged from an uninterrupted run at its first step. AtomisticTask now stores the torch, CUDA, numpy and python RNG states in each checkpoint and restores them in on_load_checkpoint, which Lightning calls before it creates the train dataloader iterator that draws the shuffle seed. They are only restored when a trainer is fitting, so loading a task with load_from_checkpoint leaves the global RNG alone. With deterministic=True, runs on QM9 that were stopped cleanly, killed or sent SIGTERM and then resumed are bitwise identical to an uninterrupted run, including weights, optimizer state and EMA. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_log_with_rank passed its parts to logging.debug like arguments of print, so with debug logging on, every call printed a TypeError traceback instead of the message. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`torch.jit.script` is deprecated and the decorator made every `import schnetpack` print a FutureWarning. A model is still scriptable as a whole, since TorchScript compiles scatter_add recursively, and full PaiNN steps show no measurable slowdown. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`run.path` is already absolute, so prefixing it with the working directory created the Aim repo at a doubled path below the working directory. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Unlike TensorBoard, CSVLogger drops the earlier rows when it logs to an existing version, so a resumed run has to keep its own version_N. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
It still targets Python 3.8 and black, and nothing refers to it. The CI workflow and pre-commit config replaced it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Models with a CastTo64 or CastMap postprocessor, as in the experiment configs, could not be converted to TorchScript, since casting looks up the dtype in a global dict. The calculator now removes these postprocessors before scripting, as spkdeploy does; postprocessing is switched off right afterwards anyway. Model files from 2.0.x still failed, because load_model gave them a plain list as electronic_embeddings. It is now an empty ModuleList, as in new models. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With load_properties left at its default of None, every sample reopened the database to read the metadata, at about 0.55 ms per read. The property names are now read once and kept until the metadata is written again. available_properties returns a copy, so that callers cannot change the cached list and with it the validation of load_properties. Based on #792. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Oct 30, 2026
Closed
fixed torch nbh list
With a single dataloader, Lightning 2.x passes the predictions of all batches as a flat list instead of one list per dataloader. PredictionWriter.write_on_epoch_end iterated over predictions[0], the dict of the first batch, and failed with "'str' object has no attribute 'items'". spkpredict therefore crashed with its default write_interval "epoch" and wrote no predictions.pt. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The callback loads the EMA average for validation and restored the raw training weights only at the start of each training epoch. Since torch_ema's restore() does not clear the stored weights: - with check_val_every_n_epoch > 1, an epoch without validation restored the weights stored at the last validation, discarding all training since then; - with validation within an epoch (val_check_interval < 1 or a step count), training continued from the EMA average after every validation. The callback now tracks whether the average is loaded and restores the raw weights at the next training batch. The flag is saved with the callback state; checkpoints without it load as before. Checkpoints and the best model still hold the EMA average. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
For atoms outside the cell, matscipy also returns cell shifts along non-periodic directions, which became offsets to periodic images that do not exist. An ethanol molecule with pbc=False outside its 6 A box got distances of up to 85 A instead of 4 A. These shifts are now set to zero. For a cell without volume, the neighbor list replaced the whole cell by a box around the atoms, which discards the lattice vectors of periodic directions: a Cu(111) slab with a zero c-vector lost 82 of its 360 pairs. This case now raises a ValueError. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
FPBE0 was "FMBD", the name of FMBD, so both collapsed into one entry of the property tables: the PBE0 forces (pbe0FOR) were never written to the database, and loading FPBE0 returned the MBD forces (vdwFOR). The total PBE0+MBD forces (QM7X.forces) were not affected. Databases built before this change hold no PBE0 forces and have to be rebuilt to get them. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Since ASE 3.23, Dynamics.__init__ calls atoms.__ase_optimizable__(), which the list of Atoms passed by BatchwiseDynamics does not have, so ASEBatchwiseLBFGS could not be created with any ASE version allowed by pyproject.toml. BatchwiseDynamics now sets the attributes of Dynamics that its run loop uses itself. Also: - BatchwiseCalculator and BatchwiseEnsembleCalculator load models with load_model, which converts models saved with older versions. - run() and irun() without steps run until converged, as in ASE, instead of not taking a single step. - The how-to notebook is enabled again. It now relaxes a batch of displaced ethanol structures and passes the fixed-atoms mask. - First tests for the module. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The previous neighbor list was reused if any of the nine cell entries and any pbc flag were unchanged, instead of all of them. A strained cubic cell keeps its zero off-diagonal entries, so the old list and its offsets, computed for the old cell, were reused: after a 2% strain of Si, distances were off by up to 0.16 A. This affected variable-cell relaxations and NPT runs through the ASE interface. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Response listed the original property names as model outputs, but stores the results under the mapped names, so every forward pass with a renamed property failed with a KeyError. The map passed by the caller is also no longer modified. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
spkdeploy loaded models with torch.load, which skips the conversion of models saved with older SchNetPack versions, so these could not be converted to TorchScript. Both force models in tests/testdata failed with "Module 'PaiNN' has no attribute 'electronic_embeddings'". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The dataset modules used to star-import schnetpack.data, so schnetpack.datasets passed all of its names on. Since these imports were made explicit, code like spk.datasets.ASEAtomsData or spk.datasets.AtomsDataModule fails with an AttributeError. Re-export schnetpack.data from schnetpack.datasets for backward compatibility. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Neither aim nor wandb is a dependency of SchNetPack, so logger=aim and logger=wandb fail on a fresh install. Say so in both configs and in the config docs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
extract_xz wrote directly to the target file and removed it only on an Exception. After a KeyboardInterrupt, the truncated .hdf5 file stayed behind, and later runs reused it because it existed. Extract to a .part file and rename it once complete. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Final 2.x release: the non-breaking changes from
dev, backported with authors kept, plus resume and CLI fixes. The breaking dataset refactor (#781) stays ondev; only its checkpoint helper is included. The full changelog is in the release notes.Also included: pytorch_lightning>=2.6.1 (fixes workers repeating the same random stream every epoch), multi-logger fix, script_model=True in MD, seed and log_every_n_steps config keys, new CI (#799), version 2.3.0.
Evidence
QM9, 200 train / 50 val, 8 epochs, seed 42, A100, interrupted in epoch 3, compared with an uninterrupted run:
Tests: 80 passed, 3 skipped, with Lightning 2.6.1 / torch 2.5.0 and with Lightning 2.6.6 / torch 2.15.0.
Merge Danger
Door: two-way
Revertible until v2.3.0 is published. Please use a merge commit or rebase, not squash, so contributors keep their commits.
Blast Radius: moderate
Users on Lightning < 2.6.1 must upgrade. black, pre-commit and pytest are no longer runtime dependencies, and the code base was reformatted with ruff, which can conflict with open branches. Checkpoints gain an rng_states entry.
🤖 Generated with Claude Code