[Bug]: Checkpoint loading fails when model doesn't use normalization

[blaz-r]

Describe the bug

With the latest update from #2130 (Thanks for the fix @alexriedel1 😄) the loading of model that was trained without normalization metrics doesn't work.
The attributes from these lines here:
https://github.com/openvinotoolkit/anomalib/blob/2bd2842ec33c6eedb351d53cf1a1082069ff69dc/src/anomalib/models/components/base/anomaly_module.py#L180-L182

are not set if the conditions above them are not met:
https://github.com/openvinotoolkit/anomalib/blob/2bd2842ec33c6eedb351d53cf1a1082069ff69dc/src/anomalib/models/components/base/anomaly_module.py#L171-L176

Leading to AttributeError. (see attached logs, and code to reproduce)

I believe that to fix this, the assignment should be done inside the if statement. If my understanding of everything here is correct, these keys are either all in the dict or neither of them are, leading to a fix along the following lines:

 if "pred_scores_normalization_class" in state_dict:
   self.anomaly_maps_normalization_metrics = self._get_instance(state_dict, "anomaly_maps_normalization_class")
   self.box_scores_normalization_metrics = self._get_instance(state_dict, "box_scores_normalization_class")
   self.pred_scores_normalization_metrics = self._get_instance(state_dict, "pred_scores_normalization_class")

   self.normalization_metrics = MetricCollection(
       {
           "anomaly_maps": self.anomaly_maps_normalization_metrics,
           "box_scores": self.box_scores_normalization_metrics,
           "pred_scores": self.pred_scores_normalization_metrics,
        },
     )

@alexriedel1 what do you think?

Dataset

MVTec

Model

PADiM

Steps to reproduce the behavior

from anomalib.models.image import Padim
from anomalib.data.image import MVTec
from anomalib.engine import Engine
from anomalib.utils.normalization import NormalizationMethod


def reproduce():
    #### CONFIGURE THIS ####
    mvtec_path = "../datasets/MVTec"
    #####

    data = MVTec(root=mvtec_path, category="bottle", num_workers=0)
    model = Padim()
    engine = Engine(normalization=NormalizationMethod.NONE, default_root_dir="./reproduce")

    engine.fit(model, datamodule=data)

    ckpt_path = engine._cache.args["default_root_dir"] / "weights" / "lightning" / "model.ckpt"

    engine.predict(model, datamodule=data, ckpt_path=ckpt_path)


if __name__ == '__main__':
    reproduce()

OS information

OS information:

• OS: Win11
• Python version: 3.10.3
• Anomalib version: 1.2.0dev
• PyTorch version: 2.3.0
• CUDA/cuDNN version: x
• GPU models and configuration: x
• Any other relevant information: x

Expected behavior

The model can still be loaded from the checkpoint, even if normalization was not used

Screenshots

No response

Pip/GitHub

GitHub

What version/branch did you use?

1.2.0dev

Configuration YAML

/

Logs

File "C:\Code\anomalib\src\anomalib\models\components\base\anomaly_module.py", line 180, in load_state_dict
    "anomaly_maps": self.anomaly_maps_normalization_metrics,
  File "C:\Users\blazr\miniconda3\envs\anomalib_env_v1\lib\site-packages\torch\nn\modules\module.py", line 1709, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'Padim' object has no attribute 'anomaly_maps_normalization_metrics'

Code of Conduct

• I agree to follow this project's Code of Conduct