How to save a storage file that can be loaded using `torch.Storage.from_file`

siahuat0727 · August 4, 2021, 8:40am

I found the code snippet here:

pytorch/pytorch/blob/d7836172169ff6db0042e2589df658fa56c1812e/test/test_torch.py#L1844-L1853


      
          def test_from_file(self):
              def assert_with_filename(filename):
                  size = 10000
                  s1 = torch.FloatStorage.from_file(filename, True, size)
                  t1 = torch.FloatTensor(s1).copy_(torch.randn(size))
          
                  # check mapping
                  s2 = torch.FloatStorage.from_file(filename, True, size)
                  t2 = torch.FloatTensor(s2)
                  self.assertEqual(t1, t2, atol=0, rtol=0)

Also, according to the doc, torch.Storage.from_file

If shared is True the file will be created if needed.

Therefore,

To create:

# create.py
import torch

fname = 'tensor.pt'
storage = torch.LongStorage.from_file(fname, True, 40)
torch.LongTensor(storage).copy_(torch.arange(40))

To read:

# read.py
import torch

fname = 'tensor.pt'
storage = torch.LongStorage.from_file(fname, True, 40)
tensor = torch.LongTensor(storage)
print(tensor)

# Demo
$ python create.py
$ python read.py
tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17,
        18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35,
        36, 37, 38, 39])