How to save and load models in PyTorch

Table of contents

Save and load PyTorch models with state_dict (recommended)

First, state_dict stores only the essential parameters of the model (such as the weights and biases), which keeps file sizes smaller and allows for easy manipulation.Second, it offers flexibility—since state_dict is a Python dictionary, you can save not only model parameters but also optimizer states and other metadata, making it easier to resume training or fine-tune models. Additionally, state_dict simplifies the process of reloading models across different environments, as it is compatible with Python’s built-in pickle module for straightforward saving and loading. This guide will take you through the steps to save and load models in PyTorch using state_dict and explain best practices for effective model management.
Second, since state_dict is a Python dictionary, you can save not only model parameters but also optimizer states and other metadata, making it easier to resume training or fine-tune models.
Additionally, state_dict simplifies the process of reloading models across different environments, as it is compatible with Python’s built-in pickle module for straightforward saving and loading. This guide will take you through the steps to save and load models in PyTorch using state_dict and explain best practices for effective model management.

Save your model in PyTorch using state_dict

torch.save(model.state_dict(), 'save/to/path/model.pth')

Load your model in PyTorch using state_dict

model = MyModelDefinition(args)
model.load_state_dict(torch.load('load/from/path/model.pth'))

Pros

PyTorch relies on Python’s pickle module, which allows Python dictionaries to be easily picked, unpickled, updated, and restored. Using state_dict provides flexibility in managing model parameters.
Along with model parameters, you can save additional elements like optimizer states and hyperparameters as key-value pairs in the state_dict. These can be easily accessed when reloading the model.

Cons:

To load a state_dict, you’ll need the exact model definition. Without it, loading the saved parameters will not work correctly.

Gotchas:

Make sure to call model.eval() when making inferences to set layers like dropout and batch normalization to evaluation mode.
Save the model using .pt or .pth extension for consistency and compatibility.

Save and load the entire PyTorch model

Save the entire PyTorch model

torch.save(model, 'save/to/path/model.pt')

Load the entire PyTorch model

model = torch.load('load/from/path/model.pt')

Pros:

Saving the entire model requires the least amount of code, making it a quick option for saving and loading.
The API for saving and loading the full model is straightforward and easy to use.The saving and loading API is more intuitive.

Cons:

Since Python’s pickle module is used internally, the saved model is tightly bound to specific class definitions and the directory structure. This means that any refactoring or changes to the file paths can cause loading issues.
Using the full model save in a different project is challenging, as the directory structure and class paths must be identical to those used when saving.

Gotchas:

Make sure to call model.eval() before making inferences to ensure proper behavior in layers like dropout and batch normalization.
Save the model using .pt or .pth extension when saving your model to maintain consistency and readability.

Save and load your PyTorch model from a checkpoint

Save a PyTorch model checkpoint

torch.save({'epoch': EPOCH,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS}, 
	    'save/to/path/model.pth')

Load a PyTorch model checkpoint

model = MyModelDefinition(args)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load('load/from/path/model.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

Gotchas:

When resuming training, remember to call model.train() to ensure the model is set to training mode.
For inference after loading a checkpoint, call model.eval() to set the model to evaluation mode.

Use W&B Artifacts for model version control

Save a PyTorch model as a W&B Artifacts

# Import
import wandb
# Save your model.
torch.save(model.state_dict(), 'save/to/path/model.pth')
# Save as artifact for version control.
run = wandb.init(project='your-project-name')
artifact = wandb.Artifact('model', type='model')
artifact.add_file('save/to/path/model.pth')
run.log_artifact(artifact)
run.finish()

Load a PyTorch model as a W&B Artifacts

import wandb
run = wandb.init()

artifact = run.use_artifact('entity/your-project-name/model:v0', type='model')
artifact_dir = artifact.download()

run.finish()

Figure 1: Check out the live artifact dashboard here

Try Weights & Biases

Click the green "Run" button below (the first time you click Run, Replit will take approx 30-45 seconds to allocate a machine)
Follow the prompts in the terminal window (the bottom right pane below)
You can resize the terminal window (bottom right) for a larger view

Recommended Reading