diff --git a/CHANGELOG.md b/CHANGELOG.md index fa946d02..da5d64c2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -70,6 +70,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Fix `WeatherDataset.__len__` off-by-one in analysis mode (was undercounting by 1 sample), include `num_past_forcing_steps` in the forecast-mode minimum-horizon check, validate forcing-side forecast horizon when forcing is present, and use `min(n_state, n_forcing)` when both are present in analysis mode; raise `IndexError` for out-of-range indices in `WeatherDataset.__getitem__` (with Python-style negative indexing support) [\#312](https://github.com/mllam/neural-lam/pull/312) @kshirajahere +- Fix `--load` to truly restore weights only by default: instead of overriding the optimizer state inside `on_load_checkpoint` (which left epoch / scheduler / callbacks inherited from the checkpoint), `train_model.py` now reconstructs the model via `ForecasterModule.load_from_checkpoint` and calls `trainer.fit(model)` with no `ckpt_path`, so all non-weight training state starts fresh. Also rename the CLI flag from `--restore_opt` to `--load_training_state` to reflect that it controls more than just the optimizer (epoch / scheduler / callbacks too), and assert it requires `--load` [\#240](https://github.com/mllam/neural-lam/pull/240) @Mani212005 + ### Maintenance - Group the existing Neural-LAM citation papers in the README under a `### Core Neural-LAM Publications` subheading for clearer structure [\#633](https://github.com/mllam/neural-lam/pull/633) @HetaviM29 diff --git a/neural_lam/models/module.py b/neural_lam/models/module.py index 215edab8..f059059d 100644 --- a/neural_lam/models/module.py +++ b/neural_lam/models/module.py @@ -822,7 +822,3 @@ def on_load_checkpoint(self, checkpoint): ) loaded_state_dict[new_key] = loaded_state_dict[old_key] del loaded_state_dict[old_key] - - if not self.restore_opt: - opt = self.configure_optimizers() - checkpoint["optimizer_states"] = [opt.state_dict()] diff --git a/neural_lam/train_model.py b/neural_lam/train_model.py index dfb199f2..410a8aa1 100644 --- a/neural_lam/train_model.py +++ b/neural_lam/train_model.py @@ -127,9 +127,13 @@ def main(input_args=None): help="Path to load model parameters from", ) runtime_group.add_argument( - "--restore_opt", + "--load_training_state", action="store_true", - help="If optimizer state should be restored with model", + help=( + "If the full training state (optimizer, epoch, LR scheduler, " + "callbacks) should be restored from the checkpoint passed to " + "--load. Without this flag, --load is weights-only." + ), ) # Model architecture @@ -371,6 +375,12 @@ def main(input_args=None): "Evaluation (--eval) without --load: no checkpoint will be loaded.", ) + if args.load_training_state and not args.load: + parser.error( + "--load_training_state requires --load to " + "specify which checkpoint to restore training state from." + ) + # Get an (actual) random run id as a unique identifier random_run_id = random.randint(0, 9999) @@ -452,7 +462,7 @@ def main(input_args=None): datastore=datastore, loss=args.loss, lr=args.lr, - restore_opt=args.restore_opt, + restore_opt=args.load_training_state, n_example_pred=args.n_example_pred, create_gif=args.create_gif, val_steps_to_log=args.val_steps_to_log, @@ -510,7 +520,19 @@ def main(input_args=None): datamodule=data_module, ckpt_path=args.load, ) + elif args.load and not args.load_training_state: + # Weights-only restore: load model weights from the checkpoint but + # start a fresh training run (epoch/scheduler/optimizer/callbacks + # all reset). Passing `ckpt_path=args.load` to `trainer.fit` would + # also restore those, which is `--load_training_state` territory. + model = load_forecaster_module_from_checkpoint( + args.load, config=config, datastore=datastore + ) + trainer.fit(model=model, datamodule=data_module) else: + # Either no checkpoint (`args.load is None`) or full resume + # (`--load_training_state`); in both cases Lightning handles the + # state. trainer.fit(model=model, datamodule=data_module, ckpt_path=args.load) diff --git a/tests/test_cli.py b/tests/test_cli.py index a2238d68..495e14a6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -134,3 +134,21 @@ def test_unsupported_logger_raises_value_error(): with pytest.raises(ValueError, match="Unsupported logger type"): setup_training_logger(datastore, args, run_name="my-run") + + +def test_load_training_state_requires_load(): + """--load_training_state without --load must fail at the CLI argument + check, before reaching `load_config_and_datastore`. The assertion lives + in train_model.main and protects against the user expecting to resume + training state without supplying a checkpoint.""" + with patch("neural_lam.train_model.load_config_and_datastore") as mock_load: + with pytest.raises(SystemExit, match="2"): + neural_lam.train_model.main.__wrapped__( + [ + "--config_path", + "dummy.yaml", + "--load_training_state", + ] + ) + + mock_load.assert_not_called() diff --git a/tests/test_train_model_warnings.py b/tests/test_train_model_warnings.py index a0b5f92a..2536f670 100644 --- a/tests/test_train_model_warnings.py +++ b/tests/test_train_model_warnings.py @@ -24,6 +24,7 @@ def test_eval_without_load_warning(eval_val, load_val, expect_warning): mock_args.val_steps_to_log = [] mock_args.var_leads_metrics_watch = "{}" mock_args.ar_steps_eval = 10 + mock_args.load_training_state = False with patch( "neural_lam.train_model.ArgumentParser.parse_args", @@ -55,6 +56,7 @@ def test_create_gif_forwarded_to_forecaster_module(): mock_args.create_gif = True mock_args.devices = ["auto"] mock_args.model = "graph_lam" + mock_args.load_training_state = False captured_kwargs = {}