Add ckpt_path option to LightningModule.test() introduced robustness issues

[mstewart141]

🐛 Bug

Trainer.test(), passing no explicit model, now looks for the best model checkpoint. However, this can fail for models that have __init__ methods that take more than just hparams, which is common.

If I have a pl model with init (kw)args (self, hparams: Namespace, fancy_thing: FancyThing) one runs into problems outlined below.

With the new codepath, Trainer.test() calls down via this path

        # if model is not given (None), ckpt_path is given,
        # load the given checkpoint for testing
        if model is None and ckpt_path is not None:
            # ckpt_path is 'best' so load the best model
            if ckpt_path == 'best':
                ckpt_path = self.checkpoint_callback.best_model_path
            model = self.get_model().load_from_checkpoint(ckpt_path)

Which soon calls down to pytorch_lightning/core/saving.py @ _load_model_state.

This correctly collects init_args_name (here hparams, fancy_thing) but the subsequent logic does not check whether all the requisite params have been found

            if args_name == 'kwargs':
                cls_kwargs = {k: v for k, v in model_args.items() if k in init_args_name}
                kwargs.update(**cls_kwargs)
            elif args_name:
                if args_name in init_args_name:
                    kwargs.update({args_name: model_args})
            else:
                args = (model_args, ) + args

        # load the state_dict on the model automatically
        model = cls(*args, **kwargs)

This means the final line of that snippet fails because fancy_thing is not passed to cls.

Even if it were checked, though, there is not a mechanism here to recover fancy_thing. Lightning might need to e.g. pickle the model init args to rehydrate them later, or provide a user hook specifying how the rehydration here should occur. That said, there may be a simple work around I do not see immediately.

To Reproduce

Call Trainer.test() passing no model, where the implicit model has init args that are not just hparams.

Code sample

I've tried to document very clearly above the code flow that triggers this, but please @ me for any more context.

Expected behavior

Rehydration should occur properly, or, as a stopgap, an error should be raised explaining that model must be explicitly be passed to Trainer.test in this case.

Additional context

Thanks for the great library; really appreciate the team's work!

[williamFalcon]

@yukw777

[yukw777]

@williamFalcon @Borda it seems like this is related to the hparams deprecation work you guys did? I thought we pickled the init arguments though... why wouldn't fancy_thing get pickled from the example?

[yukw777]

Ahh I think I know what’s going on.. would it be OK for us to accept extra args and kwargs that’d be passed into the model init function in test()?

[yukw777]

We can close this per #2423 (comment)

[Borda]

resolved in #2681