8000 `pickle.loads` throws warning with `nn.Module` · Issue #130242 · pytorch/pytorch · GitHub
[go: up one dir, main page]

Skip to content

pickle.loads throws warning with nn.Module #130242

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

Closed
NicolasHug opened this issue Jul 8, 2024 · 2 comments
Closed

pickle.loads throws warning with nn.Module #130242

NicolasHug opened this issue Jul 8, 2024 · 2 comments
Labels
module: serialization Issues related to serialization (e.g., via pickle, or otherwise) of PyTorch objects triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Milestone

Comments

@NicolasHug
Copy link
Member
NicolasHug commented Jul 8, 2024

Pickling a nn.Module that contains tensor attributes now throws a warning:

import torch
import pickle

pickle.loads(pickle.dumps(torch.nn.Linear(10, 10)))
/home/nicolashug/.miniconda3/envs/pt/lib/python3.11/site-packages/torch/storage.py:505: 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.

This seems to be due to #129239 (CC-ing authors and reviewers: @mikaylagawarecki @albanD @malfet).

What is the recommended workaround to avoid the warning? The current warning message is unfortunately unactionable in this case because the user didn't call torch.load() directly, only pickle.loads().

Note: this is related to / causing pytorch/vision#8517

torch version: 2.5.0.dev20240708+cpu

cc @mruberry @mikaylagawarecki

@mikaylagawarecki
Copy link
Contributor

I believe the root cause is that __reduce__ for Storage

pytorch/torch/storage.py

8000

Line 1206 in 16d53cb

return (_load_from_bytes, (b.getvalue(),))

_load_from_bytes here calls torch.load

pytorch/torch/storage.py

Lines 504 to 505 in 16d53cb

def _load_from_bytes(b):
return torch.load(io.BytesIO(b))

Since torch.save does not use pickle for storages, this __reduce__ is only called when pickle is used, I think setting weights_only=False in _load_from_bytes will fix this, let me send a PR

@mikaylagawarecki mikaylagawarecki added module: serialization Issues related to serialization (e.g., via pickle, or otherwise) of PyTorch objects triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module labels Jul 8, 2024
xuhancn pushed a commit to xuhancn/pytorch that referenced this issue Jul 25, 2024
Fixes pytorch#130242

Since `torch.save` does not use pickle for storages, the `torch.load` in `_load_from_bytes` should not ever be called when `torch.load`-ing a checkpoint. Setting weights_only=False explicitly in `_load_from_bytes` to avoid the weights_only warning when using the pickle module

Pull Request resolved: pytorch#130246
Approved by: https://github.com/albanD
@atalman atalman added this to the 2.4.1 milestone Aug 15, 2024
pytorchbot pushed a commit that referenced this issue Aug 15, 2024
Fixes #130242

Since `torch.save` does not use pickle for storages, the `torch.load` in `_load_from_bytes` should not ever be called when `torch.load`-ing a checkpoint. Setting weights_only=False explicitly in `_load_from_bytes` to avoid the weights_only warning when using the pickle module

Pull Request resolved: #130246
Approved by: https://github.com/albanD

(cherry picked from commit dfd1d19)
atalman pushed a commit that referenced this issue Aug 21, 2024
Fix warning when pickle.load torch.Storage (#130246)

Fixes #130242

Since `torch.save` does not use pickle for storages, the `torch.load` in `_load_from_bytes` should not ever be called when `torch.load`-ing a checkpoint. Setting weights_only=False explicitly in `_load_from_bytes` to avoid the weights_only warning when using the pickle module

Pull Request resolved: #130246
Approved by: https://github.com/albanD

(cherry picked from commit dfd1d19)

Co-authored-by: Mikayla Gawarecki <mikaylagawarecki@gmail.com>
@malfet
Copy link
Contributor
malfet commented Sep 3, 2024

Validated the fix by running:

$ python -c "import torch;import pickle; print(torch.__version__); pickle.loads(pickle.dumps(torch.nn.Linear(10, 10)))"
2.4.0+cpu
/home/nshulga/miniconda3/envs/py310-torch/lib/python3.10/site-packages/torch/storage.py:414: 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.
  return torch.load(io.BytesIO(b))

But with 2.4.1:

$ python -c "import torch;import pickle; print(torch.__version__); pickle.loads(pickle.dumps(torch.nn.Linear(10, 10)))"
2.4.1+cpu

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
module: serialization Issues related to serialization (e.g., via pickle, or otherwise) of PyTorch objects triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module
Projects
None yet
Development

No branches or pull requests

4 participants
0