Skip to content

Commit

Permalink
Use torch.save in _StorageBase.__reduce__ (#9184)
Browse files Browse the repository at this point in the history
Summary:
Previously this used the ``.toliist`` method, which converted the
storage object into a list of Python objects, and then sent those to
pickle.  For storage objects of non-trivial size, this was very slow.

Now we reuse the logic of the ``torch.save`` function to efficiently
turn the Storage object into bytes, and send those instead.  This
reduces the semantic information (it's harder to interpret the bytes)
but should be orders of magnitude more efficient when serializing data
with the pickle protocol or with copy

For future work it would be nice to develop a mechanism to get a buffer
of bytes out of a Storage object, and use that alongside the current
``from_buffer`` method.

See #9168 for context
Closes #9184

Differential Revision: D8747794

Pulled By: soumith

fbshipit-source-id: ac598e660c043788ed1ffab3d0303812886edf79
  • Loading branch information
mrocklin authored and facebook-github-bot committed Jul 6, 2018
1 parent 7b25cbb commit eadc507
Showing 1 changed file with 9 additions and 1 deletion.
10 changes: 9 additions & 1 deletion torch/storage.py
@@ -1,3 +1,5 @@
import io

import torch
from ._utils import _type, _cuda

Expand Down Expand Up @@ -28,7 +30,9 @@ def __deepcopy__(self, memo):
return new_storage

def __reduce__(self):
return type(self), (self.tolist(),)
b = io.BytesIO()
torch.save(self, b)
return (_load_from_bytes, (b.getvalue(),))

def __sizeof__(self):
return super(_StorageBase, self).__sizeof__() + self.element_size() * self.size()
Expand Down Expand Up @@ -116,5 +120,9 @@ def _new_shared(cls, size):
return cls._new_using_fd(size)


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


_StorageBase.type = _type
_StorageBase.cuda = _cuda

0 comments on commit eadc507

Please sign in to comment.