DEV Community

Super Kai (Kazuya Ito)
Super Kai (Kazuya Ito)

Posted on

save() and load() in PyTorch

Buy Me a Coffee

save() can save an object as shown below:

*Memos:

  • save() can be used with torch but not with a tensor.
  • The 1st argument with torch is obj(Required-Type:object).
  • The 2nd argument with torch is f(Required-Type:Union[str, PathLike, BinaryIO, IO[bytes]]).
  • The 3rd argument with torch is pickle_module(Optional-Default:pickle-Type:Any).
  • The 4th argument with torch is pickle_protocol(Optional-Default:2-Type:int).
  • The 5th argument with torch is _use_new_zipfile_serialization(Optional-Default:True-Type:bool).
  • The 6th argument with torch is _disable_byteorder_record(Optional-Default:False-Type:bool).

1. Create my_project folder manually:

*Memos:

my_project
Enter fullscreen mode Exit fullscreen mode

2. Save file1(str) in my_project:

import torch
import pickle

torch.save(obj="Hello World", f="my_project/file1")
# Or
# torch.save(obj="Hello World", f="my_project/file1",
#            pickle_module=pickle,
#            pickle_protocol=2,
#            _use_new_zipfile_serialization=True,
#            _disable_byteorder_record=False)

# my_project
#  └-file1 <- Here
Enter fullscreen mode Exit fullscreen mode

3. Save file2(int), file3(float), file4(complex) and file5(bool) in my_project:

import torch

torch.save(obj=3, f="my_project/file2")
torch.save(obj=3.14, f="my_project/file3")
torch.save(obj=3.14+7.j, f="my_project/file4")
torch.save(obj=True, f="my_project/file5")

# my_project
#  |-file1
#  |-file2 <- Here
#  |-file3 <- Here
#  |-file4 <- Here
#  └-file5 <- Here
Enter fullscreen mode Exit fullscreen mode

load() can load an object as shown below:

*Memos:

  • load() can be used with torch but not with a tensor.
  • The 1st argument with torch is f(Required-Type:Union[str, PathLike, BinaryIO, IO[bytes]]).
  • The 2nd argument with torch is map_location(Optional-Default:None-Type:Optional[Union[Callable[Storage, str], Storage], device, str, Dict[str, str]]]).
  • The 3rd argument with torch is pickle_module(Optional-Default:pickle-Type:Any). *It must be None if weights_only is True but pickle is set to it implicitly.
  • There is weights_only argument with torch(Optional-Default:False-Type:Optional(bool)): *Memos:
    • Basically, True should be set to it. *A warning occurs if weights_only=False isn't explicitly set.
    • It will be True by default in the later version of PyTorch.
  • There is mmap argument with torch(Optional-Default:None-Type:Optional(bool)).
  • There is pickle_load_args argument with torch(Optional-Type:Any): *Memos:
    • A keyword argument must be used: *Memos:
    • pickle_load_args= cannot be used.
    • errors= can be used.
    • It cannot be used if weights_only is True.
    • It's only for Python 3.
  • It can load a lot more types than I show below. *complex type can be loaded if weights_only is False.

1. There are file1(str), file2(int), file3(float), file4(complex) and file5(bool) in my_project as saved above with save():

my_project
 |-file1
 |-file2
 |-file3
 |-file4
 └-file5
Enter fullscreen mode Exit fullscreen mode

2. Load file1(str) from my_project:

import torch
import pickle

torch.load(f="my_project/file1", weights_only=True)
# 'Hello World'

torch.load(f="my_project/file1", map_location=None, 
           pickle_module=None, weights_only=True,
           mmap=None)
# 'Hello World'

torch.load(f="my_project/file1", map_location=None, 
           pickle_module=pickle, weights_only=False,
           mmap=None, errors="There are errors")
# 'Hello World'
Enter fullscreen mode Exit fullscreen mode

3. Load file2(int), file3(float), file4(complex) and file5(bool) in my_project. *complex type can be loaded if weights_only is False:

import torch

torch.load(f="my_project/file2", weights_only=True)
# 3

torch.load(f="my_project/file3", weights_only=True)
# 3.14

torch.load(f="my_project/file4", weights_only=False)
# (3.14+7j)

torch.load(f="my_project/file5", weights_only=True)
# True
Enter fullscreen mode Exit fullscreen mode

Top comments (0)