Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adapt to Pytorch 2.6 security improvements #98

Open
zhong-al opened this issue Feb 16, 2025 · 2 comments · May be fixed by #100
Open

Adapt to Pytorch 2.6 security improvements #98

zhong-al opened this issue Feb 16, 2025 · 2 comments · May be fixed by #100
Assignees
Labels
bug Something isn't working

Comments

@zhong-al
Copy link
Collaborator

https://pytorch.org/blog/pytorch2-6/

Also in this release as an important security improvement measure we have changed the default value for weights_only parameter of torch.load. This is a backward compatibility-breaking change, please see this forum post for more details.

A potential solution is torch.serialization.add_safe_globals (https://pytorch.org/docs/stable/notes/serialization.html).

@zhong-al zhong-al self-assigned this Feb 16, 2025
@zhong-al zhong-al added the bug Something isn't working label Feb 16, 2025
@zhong-al
Copy link
Collaborator Author

even with

from torch.nn.modules.container import Sequential
from ultralytics import YOLO
from ultralytics.nn.tasks import DetectionModel
from ultralytics.nn.modules import Conv

torch.serialization.add_safe_globals([Sequential])
torch.serialization.add_safe_globals([DetectionModel])
torch.serialization.add_safe_globals([Conv])

I get

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, do those steps only if you trust the source of the checkpoint.
(1) In PyTorch 2.6, we changed the default value of the weights_only argument in torch.load from False to True. 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 ultralytics.nn.modules.Conv was not an allowed global by default. Please use torch.serialization.add_safe_globals([Conv]) or the torch.serialization.safe_globals([Conv]) context manager to allowlist this global if you trust this class/function.

So torch.serialization.add_safe_globals doesn't work with Conv. Will look into alternatives, otherwise downgrading torch to 2.5 should work.

@zhong-al
Copy link
Collaborator Author

A better solution seems to be upgrading ultralytics version: ultralytics/yolov5#13513

This handles torch 2.6 compatibility issues with minimal effort.

Messaged @dirtmaxim on slack to determine why ultralytics~=8.0.36 was added to requirements.txt over ultralytics>=8.0.36.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

Successfully merging a pull request may close this issue.

1 participant