Best way to save many tensors of different shapes?

Question

I would like to store thousands to millions of tensors with different shapes to disk. The goal is to use them as a time series dataset. The dataset will probably not fit into memory and I will have to load samples or ranges of samples from disk.

What is the best way to accomplish this while keeping storage and access time low?

Answer 1

The easiest way to save anything in disk is by using pickle:

import pickle
import torch

a = torch.rand(3,4,5)

# save
with open('filename.pickle', 'wb') as handle:
    pickle.dump(a, handle)

# open
with open('filename.pickle', 'rb') as handle:
    b = pickle.load(handle)

You can also save things with pytorch directly, but that is just a pytorch wrapper around pikle.

import torch
x = torch.tensor([0, 1, 2, 3, 4])
torch.save(x, 'tensor.pt')

If you want to save multiple tensors in one file, you can wrap them in a dictionary:

import torch
x = torch.tensor([0, 1, 2, 3, 4])
a = torch.rand(2,3,4,5)
b = torch.zeros(37)
torch.save({"a": a, "b":b, "x", x}, 'tensors.pt')