Fast serialization of Tensors?

username: stsievert
cooked: <p>I’d like a fast serialization for a tensor. Right now, I’m using</p>
<pre><code class="lang-python">x = torch.Tensor(...)
serialized = pickle.dumps(x)
</code></pre>
<p>The tensor is not guaranteed to live on the CPU and I want to preserve CPU–GPU bandwidth, meaning I can not use</p>
<pre><code class="lang-python">x.numpy().tobytes()
</code></pre>
<p>Is there a fast serialization method for torch Tensors?</p>
<p>Notes on research I’ve done so far:</p>
<ul>
<li>the source for <code>torch.serialization</code> and it relies on pickle.dumps (plus it seems oriented towards files, not speed).</li>
<li>blosc has <a href="https://github.com/Blosc/python-blosc/blob/2a2ebcecc46b9a461883530b46fa371ffbcc0c01/blosc/toplevel.py#L418" rel="nofollow noopener"><code>compress_ptr</code></a> which could be useful with <code>x.data_ptr</code>
</li>
<li>I’ve looked at <a href="https://arrow.apache.org/docs/python/ipc.html" rel="nofollow noopener">pyarrow</a> too, but decided to ask here first</li>
</ul>

username: stsievert
cooked: <p>I should look more into CuPy and their serializers: <a href="https://chainer.readthedocs.io/en/v2-docs-cupy/reference/serializers.html" rel="nofollow noopener">https://chainer.readthedocs.io/en/v2-docs-cupy/reference/serializers.html</a></p>

username: smth
cooked: <p>pytorch defines a custom pickler, so pickle.dump with the custom pickler is actually very fast (we go into C for serializing the storage)</p>

username: stsievert
cooked: <p>Right, I should clarify.</p>
<p>I’m comparing with NumPy serialization. The picture below times serialization for NumPy and PyTorch with <code>pickle.dumps</code> on a Macbook Pro 2015.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/3/3133916dd9c9f6dcdfc2fe9ab8be25aefd4d5243.png" data-download-href="https://discuss.pytorch.org/uploads/default/3133916dd9c9f6dcdfc2fe9ab8be25aefd4d5243" title="59 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/3/3133916dd9c9f6dcdfc2fe9ab8be25aefd4d5243_2_690x437.png" alt="59 PM" data-base62-sha1="71fSzkykNej8Mf0viizNFGQjOhR" width="690" height="437" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/3/3133916dd9c9f6dcdfc2fe9ab8be25aefd4d5243_2_690x437.png, https://discuss.pytorch.org/uploads/default/optimized/2X/3/3133916dd9c9f6dcdfc2fe9ab8be25aefd4d5243_2_1035x655.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/3/3133916dd9c9f6dcdfc2fe9ab8be25aefd4d5243.png 2x" data-dominant-color="F8F8F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">59 PM</span><span class="informations">1260×798 86.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The core of my code was</p>
<pre data-code-wrap="python"><code class="lang-python">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)
</code></pre>
<p>This is with <code>torch.__version__ == 0.3.0.post4</code>.</p>

username: mrshenli
cooked: <p><a class="mention" href="/u/stsievert">@stsievert</a>  Is the fix in <a href="https://github.com/pytorch/pytorch/pull/9184" rel="nofollow noopener">#9184</a> sufficient?</p>

username: stsievert
cooked: <p>Yup. Well, I think so. That PR was actually motivated by some of the work I did over the summer. I wouldn’t be surprised if the NumPy and PyTorch timings are equally fast now (but I’d like to see the graph).</p>

username: stsievert
cooked: <p>Here’s the same graph with PyTorch 1.0.0 and NumPy 1.16.2:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/7/745b7b351a5a8018b7e55e83baba9c34ef4be811.png" data-download-href="https://discuss.pytorch.org/uploads/default/745b7b351a5a8018b7e55e83baba9c34ef4be811" title="47%20PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/7/745b7b351a5a8018b7e55e83baba9c34ef4be811_2_690x450.png" alt="47%20PM" data-base62-sha1="gBlnPhtraiwGaH7w9nK1PRZtV2F" width="690" height="450" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/7/745b7b351a5a8018b7e55e83baba9c34ef4be811_2_690x450.png, https://discuss.pytorch.org/uploads/default/optimized/2X/7/745b7b351a5a8018b7e55e83baba9c34ef4be811_2_1035x675.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/7/745b7b351a5a8018b7e55e83baba9c34ef4be811.png 2x" data-dominant-color="F7F6F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">47%20PM</span><span class="informations">1108×724 93.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: tastyminerals
cooked: <p>It would be super nice to add <code>pyarrow</code> here.</p>