pytorch
b4e49179 - Avoid temporary buffers for tensors with torch.save. (#80404)

Commit
2 years ago
Avoid temporary buffers for tensors with torch.save. (#80404) Fix torch.save _open_zipfile_writer optimization that uses a c++ stream when `f` is a os.PathLike. This fastpath requires that we don't `open()` in python if possible, so don't do it unconditionally. Fix PyTorchStreamWriter construction binding that takes a buffer object. Use py::memoryview instead of py::bytes as the former doesn't copy the data. Validated with a trivial benchmark that calls torch.save in a loop 20x with a 10M elements float32 tensor either on cpu or cuda. Saved to /dev/null. Tried two variants 'str' and 'open' In 'str' we pass the string "/dev/null" to torch.save. In 'open' we pass `open("/dev/null", "wb")` to torch.save. Timing in seconds. Before this patch: str-cpu :: 0.757 open-cpu :: 0.757 str-cuda :: 1.367 open-cuda :: 1.366 After this patch: str-cpu :: 0.256 open-cpu :: 0.251 str-cuda :: 0.896 open-cuda :: 0.834 Fixes #ISSUE_NUMBER Pull Request resolved: https://github.com/pytorch/pytorch/pull/80404 Approved by: https://github.com/jamesr66a
Author
Rodrigo Kumpera
Committer
Parents
Loading