Multiple torch.load in one file

[mariogeiger]

🐛 Bug

torch.load fails to load multiple content in one file.

To Reproduce

with open('test.pkl', 'wb') as f:
    torch.save(4, f)
    torch.save(5, f)

with open('test.pkl', 'rb') as f:
    assert torch.load(f) == 4
    assert torch.load(f) == 5

Crashes with the following error

~/anaconda3/envs/env38/lib/python3.8/site-packages/torch/serialization.py in load(f, map_location, pickle_module, **pickle_load_args)
    583                     return torch.jit.load(opened_file)
    584                 return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
--> 585         return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
    586 
    587 

~/anaconda3/envs/env38/lib/python3.8/site-packages/torch/serialization.py in _legacy_load(f, map_location, pickle_module, **pickle_load_args)
    753             "functionality.".format(type(f)))
    754 
--> 755     magic_number = pickle_module.load(f, **pickle_load_args)
    756     if magic_number != MAGIC_NUMBER:
    757         raise RuntimeError("Invalid magic number; corrupt file?")

UnpicklingError: A load persistent id instruction was encountered,
but no persistent_load function was specified.

Environment

PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (Debian 8.3.0-6) 8.3.0
Clang version: Could not collect
CMake version: version 3.13.4

Python version: 3.8 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: 9.2.148
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip] numpy==1.18.1
[pip] torch==1.6.0
[pip] torch-cluster==1.5.4
[pip] torch-geometric==1.4.3
[pip] torch-scatter==2.0.5
[pip] torch-sparse==0.6.7
[pip] torch-spline-conv==1.2.0
[pip] torchvision==0.7.0
[conda] blas                      1.0                         mkl  
[conda] cpuonly                   1.0                           0    pytorch
[conda] mkl                       2020.0                      166  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.0.15           py38ha843d7b_0  
[conda] mkl_random                1.1.0            py38h962f231_0  
[conda] numpy                     1.18.1           py38h4f9e942_0  
[conda] numpy-base                1.18.1           py38hde5b4d6_1  
[conda] pytorch                   1.6.0               py3.8_cpu_0  [cpuonly]  pytorch
[conda] torch-cluster             1.5.4                    pypi_0    pypi
[conda] torch-geometric           1.4.3                    pypi_0    pypi
[conda] torch-scatter             2.0.5                    pypi_0    pypi
[conda] torch-sparse              0.6.7                    pypi_0    pypi
[conda] torch-spline-conv         1.2.0                    pypi_0    pypi
[conda] torchvision               0.7.0                  py38_cpu  [cpuonly]  pytorch

cc @ezyang @gchanan @zou3519

[albanD]

From the forum comment: this used to work in 1.5.

[gchanan]

Is this due to the format? I.e. if you use _use_new_zipfile_serialization=False does it work?

[gchanan]

ya, it works if you pass that flag, i.e.:

import torch

with open('test.pkl', 'wb') as f:
    torch.save(4, f, _use_new_zipfile_serialization=False)
    torch.save(5, f, _use_new_zipfile_serialization=False)

with open('test.pkl', 'rb') as f:
    assert torch.load(f) == 4
    assert torch.load(f) == 5

[zou3519]

cc @jamesr66a for the new zipfile serialization -- do you know what is up or who is maintaining the zipfile serialization?

[jamesr66a]

I took a look at this. The reason multiple-loads does not work with zipfile is that the zipfile reader does not advance the position in the file (f in the example) to the end of the zip file it just read. It's possible to get this behavior working, but unfortunately it's non-trivial and likely inefficient due to the way miniz` is writing out files in the archive (i.e. file sizes aren't populated in their header, rather in a trailing data descriptor struct, so traversing to the end of the archive will likely involve scanning every byte in the file.

@mariogeiger can you instead use a different way of commingling save files? For example you could wrap the multiple saves up into another zip file using the Python zipfile library

[mariogeiger]

@jamesr66a All right, I will look for an alternative

[mariogeiger]

Alternatives are not obvious
I don't want to create temporary files every time I want to read/write in a zipfile 😞

[jamesr66a]

@mariogeiger you can use io.BytesIO to avoid writing temp files, e.g.:

import torch
import io
import zipfile

to_write = [4, 5]

with zipfile.ZipFile('foo.zip', 'w') as zf:
  for i, data in enumerate(to_write):
    f = io.BytesIO()
    torch.save(data, f)
    zf.writestr(f'{i}.pkl', f.getbuffer())

with zipfile.ZipFile('foo.zip', 'r') as zf:
  for file, expected in zip(zf.namelist(), to_write):
      with zf.open(file, 'r') as f:
          assert torch.load(io.BytesIO(f.read())) == expected

[mariogeiger]

@jamesr66a Thanks!