automl / ifBO

In-context Bayesian Optimization
MIT License
10 stars 2 forks source link

`torch.load` deprecation #12

Open eddiebergman opened 1 week ago

eddiebergman commented 1 week ago

Right now using ifbo, we get very many annoying errors relating to loading the model within ifbo/surrogate.py::__init__.py

/home/skantify/code/neps-cli/.venv/lib/python3.10/site-packages/ifbo/surrogate.py:88: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  self.model = torch.load(

I tried setting this to torch.load(..., weights_only=True), however i then get the following error. Is it possible to export a weight only model, seems it wasnt serialized as you would other models. Unfortunatly I do not have much experience with this so can not recommend the fix.

Traceback (most recent call last):
  File "/home/skantify/code/neps-cli/neps_examples/efficiency/freeze_thaw.py", line 158, in <module>
    neps.run(
  File "/home/skantify/code/neps-cli/neps/api.py", line 198, in run
    ) = _run_args(
  File "/home/skantify/code/neps-cli/neps/api.py", line 423, in _run_args
    searcher_instance = instance_from_map(
  File "/home/skantify/code/neps-cli/neps/optimizers/multi_fidelity/ifbo.py", line 140, in __init__
    self.model_policy = model_policy(
  File "/home/skantify/code/neps-cli/neps/optimizers/multi_fidelity/mf_bo.py", line 237, in __init__
    super().__init__(*args, **kwargs)
  File "/home/skantify/code/neps-cli/neps/optimizers/multi_fidelity/mf_bo.py", line 198, in __init__
    self.surrogate_model = instance_from_map(
  File "/home/skantify/code/neps-cli/neps/utils/common.py", line 363, in instance_from_map
    instance = instance()  # type: ignore
  File "/home/skantify/code/neps-cli/neps/optimizers/bayesian_optimization/models/ftpfn.py", line 69, in __init__
    self.ftpfn = FTPFN(target_path=target_path, version=version)
  File "/home/skantify/code/neps-cli/.venv/lib/python3.10/site-packages/ifbo/surrogate.py", line 89, in __init__
    self.model = torch.load(
  File "/home/skantify/code/neps-cli/.venv/lib/python3.10/site-packages/torch/serialization.py", line 1096, in load
    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
_pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options 
    (1) Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
    (2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
    WeightsUnpickler error: Unsupported global: GLOBAL ifbo.transformer.TransformerModel was not an allowed global by default. Please use `torch.serialization.add_safe_globals([TransformerModel])` to allowlist this global if you trust this class/function.
eddiebergman commented 1 week ago

Just the code snippet that specifies weight_only that will cause the above crash:

        self.model = torch.load(
            os.path.join(self.target_path, WEIGHTS_FINAL_NAME(version)),
            map_location=self.device if self.device is not None else torch.device("cpu"),
            weights_only=True,
        )
Neeratyoy commented 1 week ago

@herilalaina is this fixable or would require retraining of our surrogate?

herilalaina commented 6 days ago

I would stick with weights_only=False. It allows to store and load both the weights (tensors) and the model specification (see here) from a single checkpoint file. Otherwise, we would need to manually provide the model’s hyperparameters when loading.