[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
pytorch
54b6188c - [fix] allow saving python attr on Tensor and Parameter via torch.save (#81616)
