test_torch_serde.py

import time
import io
import msgpack
import pickle
import torch

x = torch.rand(100, 100)
n_runs = 30

# raw torch bytes
start_time = time.time()
size = 0
for _ in range(n_runs):
    for row in x:
        buffer = io.BytesIO()
        torch.save(row, buffer)
        size += buffer.tell()
        buffer.seek(0)
        new_x = torch.load(buffer)
end_time = time.time()
print(f"Raw torch bytes: {end_time - start_time} seconds")
print(f"Size: {size} bytes")

# pickle
start_time = time.time()
size = 0
for _ in range(n_runs):
    for row in x:
        data = pickle.dumps(row)
        size += len(data)
        new_x = pickle.loads(data)
end_time = time.time()
print(f"Pickle: {end_time - start_time} seconds")
print(f"Size: {size} bytes")

# msgpack (row-wise tolist)
start_time = time.time()
size = 0
for _ in range(n_runs):
    for row in x:
        data = msgpack.packb(row.tolist())
        size += len(data)
        new_x = torch.tensor(msgpack.unpackb(data))
end_time = time.time()
print(f"Msgpack: {end_time - start_time} seconds")
print(f"Size: {size} bytes")

# msgpack (batch-wise tolist)
start_time = time.time()
size = 0
for _ in range(n_runs):
    x_list = x.tolist()
    for row in x_list:
        data = msgpack.packb(row)
        size += len(data)
        new_x = torch.tensor(msgpack.unpackb(data))
end_time = time.time()
print(f"Msgpack (batch-wise tolist): {end_time - start_time} seconds")
print(f"Size: {size} bytes")