kshitij12345
f74946324e
[fix] allow saving python attr on Tensor and Parameter via torch.save ( #81616 )
...
Fixes: https://github.com/pytorch/pytorch/issues/72129
TODO:
* [x] Fix for Parameter
Benchmark
(Measurable diff for small tensors)
```
[-------------- Save and Load --------------]
| After PR | Before PR
1 threads: ----------------------------------
() | 111.7 | 106.9
(4, 4) | 114.4 | 109.2
(128, 128) | 135.2 | 128.3
(1024, 1024) | 1431.9 | 1431.3
Times are in microseconds (us).
```
<details>
<summary> Benchmark Script </summary>
```python
import torch
from torch.testing._internal.common_utils import BytesIOContext
from torch.utils import benchmark
import pickle
shapes = ((), (4, 4), (128, 128), (1024, 1024))
sizes = [1, 64, 1024, 10000]
results = []
def save_load_fn(t):
with BytesIOContext() as f:
torch.save(t, f)
f.seek(0)
torch.load(f)
for shape in shapes:
t = torch.randn(shape)
label = 'Save and Load'
sub_label = f'{shape}'
results.append(benchmark.Timer(
stmt='save_load_fn(t)',
globals={'t': t, 'save_load_fn':save_load_fn},
label=label,
sub_label=sub_label,
description='Before PR',
).blocked_autorange(min_run_time=2))
compare = benchmark.Compare(results)
compare.print()
with open('before_pr.pkl', 'wb') as f:
pickle.dump(results, f)
# with open('after_pr.pkl', 'rb') as f:
# after_pr = pickle.load(f)
# with open('before_pr.pkl', 'rb') as f:
# before_pr = pickle.load(f)
# compare = benchmark.Compare(after_pr + before_pr)
# compare.print()
```
</details>
NOTE : **BC-Breaking** : After this PR, all tensors (also regular tensors) will be serialised using `_rebuild_from_type_v2`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81616
Approved by: https://github.com/albanD , https://github.com/kurtamohler
2022-11-11 21:11:12 +00:00
Kshiteej K
54b6188cc6
[fix] allow saving python attr on Tensor and Parameter via torch.save ( #81616 )
...
Fixes: https://github.com/pytorch/pytorch/issues/72129
TODO:
* [x] Fix for Parameter
Benchmark
(Measurable diff for small tensors)
```
[-------------- Save and Load --------------]
| After PR | Before PR
1 threads: ----------------------------------
() | 111.7 | 106.9
(4, 4) | 114.4 | 109.2
(128, 128) | 135.2 | 128.3
(1024, 1024) | 1431.9 | 1431.3
Times are in microseconds (us).
```
<details>
<summary> Benchmark Script </summary>
```python
import torch
from torch.testing._internal.common_utils import BytesIOContext
from torch.utils import benchmark
import pickle
shapes = ((), (4, 4), (128, 128), (1024, 1024))
sizes = [1, 64, 1024, 10000]
results = []
def save_load_fn(t):
with BytesIOContext() as f:
torch.save(t, f)
f.seek(0)
torch.load(f)
for shape in shapes:
t = torch.randn(shape)
label = 'Save and Load'
sub_label = f'{shape}'
results.append(benchmark.Timer(
stmt='save_load_fn(t)',
globals={'t': t, 'save_load_fn':save_load_fn},
label=label,
sub_label=sub_label,
description='Before PR',
).blocked_autorange(min_run_time=2))
compare = benchmark.Compare(results)
compare.print()
with open('before_pr.pkl', 'wb') as f:
pickle.dump(results, f)
# with open('after_pr.pkl', 'rb') as f:
# after_pr = pickle.load(f)
# with open('before_pr.pkl', 'rb') as f:
# before_pr = pickle.load(f)
# compare = benchmark.Compare(after_pr + before_pr)
# compare.print()
```
</details>
NOTE : **BC-Breaking** : After this PR, all tensors (also regular tensors) will be serialised using `_rebuild_from_type_v2`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81616
Approved by: https://github.com/albanD , https://github.com/kurtamohler
2022-11-03 09:57:47 +00:00
Nikita Shulga
961ebca225
Add weights_only option to torch.load ( #86812 )
...
This addresses the security issue in default Python's `unpickler` that allows arbitrary code execution while unpickling.
Restrict classes allowed to be unpicked to in `None`, `int`, `bool`, `str`, `float`, `list`, `tuple`, `dict`/`OrderedDict` as well as `torch.Size`, `torch.nn.Param` as well as `torch.Tensor` and `torch.Storage` variants.
Defaults `weights_only` is set to `False`, but allows global override to safe only load via `TORCH_FORCE_WEIGHTS_ONLY_LOAD` environment variable.
To some extent, addresses https://github.com/pytorch/pytorch/issues/52596
Pull Request resolved: https://github.com/pytorch/pytorch/pull/86812
Approved by: https://github.com/ezyang
2022-10-21 01:09:50 +00:00
