fix lightning hparams conflict

[AntonioMirarchi]

DO NOT MERGE YET!

In PyTorch Lightning, the LightningModule and LightningDataModule are two separate components with distinct responsibilities. Occasionally, conflicts in hyperparameters may arise between the model and the data loading process, causing issues during training or inference. For example, in the case that I'm reporting the error comes out when the test is going to be performed because the best_model.ckpt was saved before than I resumed a training using load_model and from this the keys mismatch. One approach to fix it, is to update the data.hparams using the model.hparams. (change proposed by this PR).

It's also true that maybe should not be allowed to have a lot of difference in the hparams between the data and the model but just some keys like load_model; if this is the case maybe the best solution is to write a function like:

def fix_lightning_conflict(model, data):
    """This function inspects the hyperparameters of the model and data and fix any conflicts by providing a consistent 
        set of hyperparameters for both components. It ensures that key hyperparameters such as batch size, learning rate, etc., 
        are coherent between the model and data loading only if the hparms key is in the allowed_conflicts list"""

    allowed_conflicts = ["load_model"]
    model_hparams = model.hparams
    data_hparams = data.hparams

    for parm in model_hparams:
        if parm in allowed_conflicts and model_hparams[parm] != data_hparams[parm]:
            rank_zero_warn(
                f"Conflict between model and dataset hparams for {parm}! "
                f"Using {model_hparams[parm]} from the model."
            )
            data_hparams[parm] = model_hparams[parm]

    return data_hparams

[AntonioMirarchi]

I'm not sure if there's a more efficient way to improve this implementation, if you have it just propose it. Moreover, if you believe that addressing this error, which arises occasionally, is not worth the effort, we can consider closing this pull request (PR).

[RaulPPelaez]

I am of the opinion that if the only difference in hparams is load_model (why is that an "hparam" anyway) then its ok. Maybe you can simply make: data.hparams.load_model = model.hparams.load_model Or similarly for a list of allowed discrepancies.

[AntonioMirarchi]

It works like a dict so you can do: data.hparams['load_model']= model.hparams['load_model']

[AntonioMirarchi]

It's possible to use the hparams_file argument in LNNP.load_from_checkpoint() to directly specify the path to the YAML file containing the hyperparameters used for training. This ensures that the model (used to test) has the same hyperparameters as the DataModule. However, it's important to note that this approach updates all hyperparameters without distinguishing between hyperparameters and parameters. Alternatively, it's possible to modify the save_hyperparameters called in the DataModule to focus on specific parameters rather than including all the arguments from the input YAML file.

I'm actually training with the version present on this PR.

[stefdoerr]

I feel like it's safer to just update those two parameters than to load the whole config in.

[AntonioMirarchi]

The dtype arg is not a problem anymore, we fixed it in #208. The original idea was to update just some parameters, essentially the ones involved in the no-matching error. The problem is that the collision check on hparams between LightningModule and DataModule (check this for more) works with the hparams_initial and I was not able to modify this dict in any way (AFAIK it works differently from the hparams.