Serialization of data within a tensor is slow

[mrocklin]

Issue description

When naively serializing a pytorch tensor (or model) with pickle the process takes longer than it probably should when comparing it to Numpy serialization of the same data. It appears that at some point we're converting things into a list rather than returning the buffers directly.

Code example

In [1]: import numpy, torch, pickle

In [2]: x = numpy.random.random((10000, 10000))

In [3]: t = torch.tensor(x)

In [4]: %time len(pickle.dumps(x)) / 1e6                    # around 1GB/s
CPU times: user 298 ms, sys: 415 ms, total: 713 ms
Wall time: 711 ms
Out[4]: 800.000162

In [5]: %time len(pickle.dumps(t)) / 1e6                    # around 50MB/s
CPU times: user 14.6 s, sys: 1.03 s, total: 15.7 s
Wall time: 15.7 s
Out[5]: 900.200098

The majority of this time is spent in converting the t.storage() object into a list

In [11]: %time _ = t.storage().tolist()
CPU times: user 12.3 s, sys: 891 ms, total: 13.2 s
Wall time: 13.2 s

Instead, we might consider passing around a numpy array, buffer, or memoryview, each of which will serialize much more quickly than converting to many Python objects

In [1]: import torch

In [2]: torch.__version__
Out[2]: '0.4.0'

[soumith]

the standard pickler will be slow. if you use torch.save to a file or file-like object, it'll be much faster, as it goes through our custom pickling logic.

Have a look here: https://github.com/pytorch/pytorch/blob/master/torch/serialization.py#L212-L286

[mrocklin]

Right, I'm suggesting making Torch's implementation of pickle fast.

[mrocklin]

Separately, what is the most efficient way to convert a torch model/tensor into a set of bytes? Pass an io.BytesIO to the save function?

[soumith]

@mrocklin the fastest way is to do t.numpy() and use whatever you'd do to numpy arrays. We dont natively provide conversion to the PyBuffer interface.

[soumith]

t.numpy() is a free operation, no memcpy, not much going on except to setup some C structs.

[mrocklin]

I understand. Would the PyTorch community accept a PR that uses numpy within the __reduce__ methods in order to improve serialization performance of tensor and model objects with naive use of pickle?

[mrocklin]

For others, looks like using io.BytesIO gets up to about 1GB/s.

In [1]: from torchvision.models.resnet import resnet18
   ...: model = resnet18(pretrained=True)
   ...: 
   ...: 

In [2]: import torch

In [3]: import io

In [4]: bio = io.BytesIO()

In [5]: %%time
   ...: torch.save(model, bio)
   ...: b = bio.getvalue()
   ...: 
CPU times: user 32 ms, sys: 16.4 ms, total: 48.4 ms
Wall time: 47.7 ms

In [6]: len(b) / 0.047 / 1e6 # MB/s
Out[6]: 996.619

And then we can reconstitute

io = io.BytesIO(b)
model2 = torch.load(io)

[mrocklin]

For context, I maintain a parallel computing library, Dask, and users are naively passing around PyTorch objects and getting poor performance. I can special-case PyTorch models to use the trick above, but it might be a good idea to make PyTorch's general serialization solution decently fast for other libraries that run into the same problem. I think that this is probably pretty easy to do.

dask/dask-ml#281 (comment)

[soumith]

Would the PyTorch community accept a PR that uses numpy within the reduce methods in order to improve serialization performance

I'll discuss with the team and get back to you in a couple of days. We've avoided a dependence on numpy for functionality so far, but it's been a while since we discussed this.

I can special-case PyTorch models to use the trick above

for the moment, this seems like a good idea.

[mrocklin]

Alternatively is there a nice way to get a memoryview out of a storage object?

…

On Wed, Jul 4, 2018, 2:45 PM Soumith Chintala ***@***.***> wrote: Would the PyTorch community accept a PR that uses numpy within the *reduce* methods in order to improve serialization performance I'll discuss with the team and get back to you in a couple of days. We've avoided a dependence on numpy for functionality so far, but it's been a while since we discussed this. I can special-case PyTorch models to use the trick above for the moment, this seems like a good idea. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#9168 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AASszNVqXOMxYD6Qao-TLr7RoxIvJm-Vks5uDQ0_gaJpZM4VCxH2> .

[mrocklin]

This would be the pure python native solution

…

On Wed, Jul 4, 2018, 2:54 PM Matthew Rocklin ***@***.***> wrote: Alternatively is there a nice way to get a memoryview out of a storage object? On Wed, Jul 4, 2018, 2:45 PM Soumith Chintala ***@***.***> wrote: > Would the PyTorch community accept a PR that uses numpy within the > *reduce* methods in order to improve serialization performance > > I'll discuss with the team and get back to you in a couple of days. We've > avoided a dependence on numpy for functionality so far, but it's been a > while since we discussed this. > > I can special-case PyTorch models to use the trick above > > for the moment, this seems like a good idea. > > — > You are receiving this because you were mentioned. > Reply to this email directly, view it on GitHub > <#9168 (comment)>, > or mute the thread > <https://github.com/notifications/unsubscribe-auth/AASszNVqXOMxYD6Qao-TLr7RoxIvJm-Vks5uDQ0_gaJpZM4VCxH2> > . >

[soumith]

MemoryView requires us to implement the Py_buffer interface, which we haven't. Implementing Py_buffer interface across Py2 and Py3 is really complicated (notes from last time we tried doing it).

[mrocklin]

Would you welcome a PR implementing the buffer protocol?

…

On Wed, Jul 4, 2018, 3:00 PM Soumith Chintala ***@***.***> wrote: MemoryView requires us to implement the Py_buffer interface, which we haven't. Implementing Py_buffer interface across Py2 and Py3 is really complicated (notes from last time we tried doing it). — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#9168 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AASszLPxJIswjHCgw8X5wHrqvAJxcDdYks5uDRDGgaJpZM4VCxH2> .

[mrocklin]

Alternatively, PyTorch clearly has code to turn storage objects efficiently into bytestreams. It must do this with torch.save. Is there an internal method somewhere to turn a storage object directly into bytes and then back?

For example, a fully usable solution for pickle would be to call the torch.save code in the comment above and just return those bytes. This isn't quite as clean, but would behave well and doesn't require much work.