Fast serialization of Tensors?


(Scott) #1

I’d like a fast serialization for a tensor. Right now, I’m using

x = torch.Tensor(...)
serialized = pickle.dumps(x)

The tensor is not guaranteed to live on the CPU and I want to preserve CPU–GPU bandwidth, meaning I can not use

x.numpy().tobytes()

Is there a fast serialization method for torch Tensors?

Notes on research I’ve done so far:

  • the source for torch.serialization and it relies on pickle.dumps (plus it seems oriented towards files, not speed).
  • blosc has compress_ptr which could be useful with x.data_ptr
  • I’ve looked at pyarrow too, but decided to ask here first

(Scott) #2

I should look more into CuPy and their serializers: https://chainer.readthedocs.io/en/v2-docs-cupy/reference/serializers.html


#3

pytorch defines a custom pickler, so pickle.dump with the custom pickler is actually very fast (we go into C for serializing the storage)


(Scott) #4

Right, I should clarify.

I’m comparing with NumPy serialization. The picture below times serialization for NumPy and PyTorch with pickle.dumps on a Macbook Pro 2015.

The core of my code was

def stat(x, serialize=pickle.dumps):
    start = time.time()
    msg = serialize(x)
    return {'time': time.time() - start, 'bytes': len(msg)}

# ... other functions, for-loops, etc
x = np.random.randn(n).astype('float32')
y = torch.Tensor(x)

This is with torch.__version__ == 0.3.0.post4.