Best way to save many tensors of different shapes?

user_3666302 0
Viewed 7165

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?

pytorch io
2 Answers

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')
user_46144040

h5py lets you save lots of tensors into the same file, and you don't have to be able to fit the entire file contents into memory. h5py will store tensors directly to disk, and you can load tensors you want when you want. It allows slicing of these tensors, at load and save time, which works in a similar way, i.e. no need to load entire tensor into memory, in order to load a slice of it, or in order to save a slice of it.

user_2127310
Related