DEV Community

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

Posted on • Edited on

save() and load() in PyTorch

Buy Me a Coffee

*Memos:

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).
  • Basically, it's used to save a model.

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:
    • weights_only= must be used.
    • 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)). *mmap= must be used.
  • 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.
  • Basically, it's used to save a model.

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

AWS Security LIVE!

Join us for AWS Security LIVE!

Discover the future of cloud security. Tune in live for trends, tips, and solutions from AWS and AWS Partners.

Learn More

Top comments (0)

Billboard image

The Next Generation Developer Platform

Coherence is the first Platform-as-a-Service you can control. Unlike "black-box" platforms that are opinionated about the infra you can deploy, Coherence is powered by CNC, the open-source IaC framework, which offers limitless customization.

Learn more

👋 Kindness is contagious

Explore a sea of insights with this enlightening post, highly esteemed within the nurturing DEV Community. Coders of all stripes are invited to participate and contribute to our shared knowledge.

Expressing gratitude with a simple "thank you" can make a big impact. Leave your thanks in the comments!

On DEV, exchanging ideas smooths our way and strengthens our community bonds. Found this useful? A quick note of thanks to the author can mean a lot.

Okay