How to convert a torch tensor into a byte string?

Question

I'm trying to serialize a torch tensor using protobuf and it seems using BytesIO along with torch.save() doesn't work. I have tried:

import torch 
import io
x = torch.randn(size=(1,20))
buff = io.BytesIO()
torch.save(x, buff)
print(f'buffer: {buff.read()}')

to no avail as it results in b'' in the output? How should I be going about this?

Answer 1

You need to seek to the beginning of the buffer before reading:

import torch 
import io
x = torch.randn(size=(1,20))
buff = io.BytesIO()
torch.save(x, buff)
buff.seek(0)  # <--  this is what you were missing
print(f'buffer: {buff.read()}')

gives you this magnificent output:

 buffer: b'PK\\x03\\x04\\x00\\x00\\x08\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x10\\x00\\x12\\x00archive/data.pklFB\\x0e\\x00ZZZZZZZZZZZZZZ\\x80\\x02ctorch._utils\\n_rebuild_tensor_v2\\nq\\x00((X\\x07\\x00\\x00\\x00storageq\\x01ctorch\\nFloatStorage\\nq\\x02X\\x0f\\x00\\x00\\x00140417054790352q\\x03X\\x03\\x00\\x00\\x00cpuq\\x04K\\x14tq\\x05QK\\x00K\\x01K\\x14\\x86q\\x06K\\x14K\\x01\\x86q\\x07\\x89ccollections\\nOrderedDict\\nq\\x08)Rq\\ttq\\nRq\\x0b.PK\\x07\\x08\\xf3\\x08u\\x13\\xa8\\x00\\x00\\x00\\xa8\\x00\\x00\\x00PK\\x03\\x04\\x00\\x00\\x08\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1c\\x00\\x0e\\x00archive/data/140417054790352FB\\n\\x00ZZZZZZZZZZ\\xba\\xf3x?\\xb5\\xe2\\xc4=)R\\x89\\xbfM\\x08\\x19\\xbfo%Y\\xbf\\x05\\xc0_\\xbf\\x03N4\\xbe\\xdd_ \\xc0&\\xc4\\xb5?\\xa7\\xfd\\xc4?f\\xf1$?Ll\\xa6?\\xee\\x8e\\x80\\xbf\\x88Uq?.<\\xd8?{\\x08\\xb2?\\xb3\\xa3\\xba>q\\xcd\\xbc?\\xba\\xe3h\\xbd\\xcan\\x11\\xc0PK\\x07\\x08A\\xf3\\xdc>P\\x00\\x00\\x00P\\x00\\x00\\x00PK\\x03\\x04\\x00\\x00\\x08\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x0f\\x003\\x00archive/versionFB/\\x00ZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZ3\\nPK\\x07\\x08\\xd1\\x9egU\\x02\\x00\\x00\\x00\\x02\\x00\\x00\\x00PK\\x01\\x02\\x00\\x00\\x00\\x00\\x08\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\xf3\\x08u\\x13\\xa8\\x00\\x00\\x00\\xa8\\x00\\x00\\x00\\x10\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00archive/data.pklPK\\x01\\x02\\x00\\x00\\x00\\x00\\x08\\x08\\x00\\x00\\x00\\x00\\x00\\x00A\\xf3\\xdc>P\\x00\\x00\\x00P\\x00\\x00\\x00\\x1c\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xf8\\x00\\x00\\x00archive/data/140417054790352PK\\x01\\x02\\x00\\x00\\x00\\x00\\x08\\x08\\x00\\x00\\x00\\x00\\x00\\x00\\xd1\\x9egU\\x02\\x00\\x00\\x00\\x02\\x00\\x00\\x00\\x0f\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xa0\\x01\\x00\\x00archive/versionPK\\x06\\x06,\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x1e\\x03-\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x03\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\xc5\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x12\\x02\\x00\\x00\\x00\\x00\\x00\\x00PK\\x06\\x07\\x00\\x00\\x00\\x00\\xd7\\x02\\x00\\x00\\x00\\x00\\x00\\x00\\x01\\x00\\x00\\x00PK\\x05\\x06\\x00\\x00\\x00\\x00\\x03\\x00\\x03\\x00\\xc5\\x00\\x00\\x00\\x12\\x02\\x00\\x00\\x00\\x00'

Answer 2

Use BytesIO.getvalue method.

Apart from seek -ing and read -ing, you can also use the getvalue method of the io.BytesIO object. It does the seek - read internally and returns the stored bytes:

In [1121]: x = torch.randn(size=(1,20))                                                                                                                        
buff = io.BytesIO()                                                                                                                                            
torch.save(x, buff)                                                                                                                                            
print(f'buffer: {buff.getvalue()}')                                                                                                                            
                                                                                                                                                               
buffer: b'PK\x03\x04\x00\x00\x08\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x10\x00\x12\x00archive/data.pklFB\x0e\\
x00ZZZZZZZZZZZZZZ\x80\x02ctorch._utils\n_rebuild_tensor_v2\nq\x00((X\x07\x00\x00\x00storageq\x01ctorch\nFloatStorage\nq\x02X\x01\x00\x00\x000q\x03X\x03\x00\x0\
0\x00cpuq\x04K\x14tq\x05QK\x00K\x01K\x14\x86q\x06K\x14K\x01\x86q\x07\x89ccollections\nOrderedDict\nq\x08)Rq\ttq\nRq\x0b.PK\x07\x08\x949f)\x9a\x00\x00\x00\x9a\\
x00\x00\x00PK\x03\x04\x00\x00\x08\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0e\x00*\x00archive/data/0FB&\x00ZZZZZZZZZZZZZZZ\
ZZZZZZZZZZZZZZZZZZZZZZZ\xff*\x1f\xbfM\xaa\x16?\x9fB\xbd?\x14\xee\xb4\xbe\xbc\x83^>l.\xba>\x8d\xc0\x1f\xbfZ\x06\x03\xbe\xe0(B\xbe^[\xf8\xbeE\x83\x9f\xbfUo\xc0\\
xbd\xbaX\xb7?\x83MH\xbf\xc0\x0c\xbb\xbf\xa4s\xc9?\x84\x8b\xd9\xbf\xa1\x91\xa1\xbf\xc6,\x0c?kxW?PK\x07\x08/\n\x02&P\x00\x00\x00P\x00\x00\x00PK\x03\x04\x00\x00\\
x08\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0f\x003\x00archive/versionFB/\x00ZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZZ\
ZZZ3\nPK\x07\x08\xd1\x9egU\x02\x00\x00\x00\x02\x00\x00\x00PK\x01\x02\x00\x00\x00\x00\x08\x08\x00\x00\x00\x00\x00\x00\x949f)\x9a\x00\x00\x00\x9a\x00\x00\x00\x1\
0\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00archive/data.pklPK\x01\x02\x00\x00\x00\x00\x08\x08\x00\x00\x00\x00\x00\x00/\n\x02&P\x00\x\
00\x00P\x00\x00\x00\x0e\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xea\x00\x00\x00archive/data/0PK\x01\x02\x00\x00\x00\x00\x08\x08\x00\x00\x00\x00\x0\
0\x00\xd1\x9egU\x02\x00\x00\x00\x02\x00\x00\x00\x0f\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xa0\x01\x00\x00archive/versionPK\x06\x06,\x00\x00\x00\\
x00\x00\x00\x00\x1e\x03-\x00\x00\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\xb7\x00\x00\x00\x00\x00\x00\x00\x\
12\x02\x00\x00\x00\x00\x00\x00PK\x06\x07\x00\x00\x00\x00\xc9\x02\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00PK\x05\x06\x00\x00\x00\x00\x03\x00\x03\x00\xb7\x00\x00\
\x00\x12\x02\x00\x00\x00\x00' 

---

getvalue also works the same for any io.StringIO object but instead of bytes, it returns the stored string, as expected.

Answer 3

Instead of using ByteIO directly you could use pickle.loads/dumps .

According to this discuss thread and the linked PR discussion pytorch's custom pickling handler does ultimately use torch.save anyways, but needs to serialize less objects, resulting in a 469 length vs 811 length bytes string. Not sure which is faster though.

import pickle, torch

x = torch.randn(size=(1,20))

pickled = pickle.dumps(x)

print(f"length of {len(pickled)} vs 811 for the BytesIO approaches")
# => length of 469 vs 811 for the BytesIO approaches

print(pickled)
# => b'\x80\x04\x95\xca\x01\x00\x00\x00\x00\x00\x00\x8c\x0ctorch._utils\x94\x8c\x12_rebuild_tensor_v2\x94\x93\x94(\x8c\rtorch.storage\x94\x8c\x10_load_from_bytes\x94\x93\x94B?\x01\x00\x00\x80\x02\x8a\nl\xfc\x9cF\xf9 j\xa8P\x19.\x80\x02M\xe9\x03.\x80\x02}q\x00(X\x10\x00\x00\x00protocol_versionq\x01M\xe9\x03X\r\x00\x00\x00little_endianq\x02\x88X\n\x00\x00\x00type_sizesq\x03}q\x04(X\x05\x00\x00\x00shortq\x05K\x02X\x03\x00\x00\x00intq\x06K\x04X\x04\x00\x00\x00longq\x07K\x04uu.\x80\x02(X\x07\x00\x00\x00storageq\x00ctorch\nFloatStorage\nq\x01X\x08\x00\x00\x0097668208q\x02X\x03\x00\x00\x00cpuq\x03K\x14Ntq\x04Q.\x80\x02]q\x00X\x08\x00\x00\x0097668208q\x01a.\x14\x00\x00\x00\x00\x00\x00\x00j\x0b\xeb?4\xcc\x1b?\xfa\xa9$\xbe\x86\xc69>\xab\xd0#?\x14}H>\xd16#\xc0\x04\x8e\x0b?\xadg\x1a?\xe2\xb5g??\xf0\x83\xbf\x02E\x89\xbf0\xf8\xe0>\x04\xc8L?\x89\xfc\x06?\xa0\xcf)\xbf\xb9\xa7\x1b@^\xd3\xa0\xbfJ\xda\xbd\xbfE\x95\x99\xbe\x94\x85\x94R\x94K\x00K\x01K\x14\x86\x94K\x14K\x01\x86\x94\x89\x8c\x0bcollections\x94\x8c\x0bOrderedDict\x94\x93\x94)R\x94t\x94R\x94.'

print(pickle.loads(pickled))
# => tensor([[ 0.7662,  0.7422,  0.1888,  ..., -0.2035,  1.0845, -0.9637]])