Loading model state_dict in PyTorch

Introduction

Loading a model's state_dict lets you reuse a saved model's learned settings. This saves time and helps continue training or make predictions.

You want to continue training a model from where you left off.

You want to use a pretrained model to make predictions on new data.

You want to share a trained model with someone else.

You want to test a model saved earlier without retraining.

You want to fine-tune a model on a new dataset.

Syntax

PyTorch

model.load_state_dict(torch.load(PATH))

model is your PyTorch model instance.

PATH is the file path to the saved state_dict.

Examples

Loads weights from 'model_weights.pth' into the model.

PyTorch

model.load_state_dict(torch.load('model_weights.pth'))

Loads the state_dict first, then applies it to the model.

PyTorch

state_dict = torch.load('weights.pth')
model.load_state_dict(state_dict)

Loads weights but allows missing or extra keys without error.

PyTorch

model.load_state_dict(torch.load('model.pth'), strict=False)

Sample Model

This program creates a simple linear model, saves its weights, changes them to zero, then reloads the saved weights to restore the original values.

PyTorch

import torch
import torch.nn as nn

# Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(2, 1)

    def forward(self, x):
        return self.linear(x)

# Create model instance and print initial weights
model = SimpleModel()
print('Initial weights:', model.linear.weight)

# Save the model's state_dict
PATH = 'simple_model.pth'
torch.save(model.state_dict(), PATH)

# Change weights manually
with torch.no_grad():
    model.linear.weight.fill_(0)
print('Weights after zeroing:', model.linear.weight)

# Load saved weights back
model.load_state_dict(torch.load(PATH))
print('Weights after loading:', model.linear.weight)

Important Notes

Always use the same model architecture when loading a saved state_dict.

Use strict=False if your saved weights don't exactly match the model.

State_dict only saves parameters, not the entire model code.

Summary

Loading a state_dict restores saved model weights.

Use torch.load(PATH) to load the saved weights.

Apply weights with model.load_state_dict().

Practice

1. What does model.load_state_dict() do in PyTorch?

easy

A. It loads saved model weights into the model.

B. It saves the current model weights to a file.

C. It initializes a new model architecture.

D. It compiles the model for training.

Solution

1. Step 1: Understand the purpose of load_state_dict

   This function is used to load previously saved weights into a model.

2. Step 2: Differentiate from other functions

   Saving weights uses state_dict() with torch.save(), not load_state_dict().

3. Final Answer:

   It loads saved model weights into the model. -> Option A

4. Quick Check:

   Load weights = load_state_dict() [OK]

Hint: Remember: load_state_dict loads weights, not saves them [OK]

Common Mistakes:

• Confusing loading weights with saving weights
• Thinking it initializes model architecture
• Assuming it compiles the model

2. Which of the following is the correct syntax to load a saved state dictionary from a file model.pth into a model named model?

easy

A. model.load_state_dict(torch.load('model.pth'))

B. model.load(torch.load_state_dict('model.pth'))

C. torch.load_state_dict(model, 'model.pth')

D. model.load_state_dict('model.pth')

Solution

1. Step 1: Identify correct function usage

   The correct way is to first load the saved weights with torch.load() and then pass them to model.load_state_dict().

2. Step 2: Check syntax correctness

   model.load_state_dict(torch.load('model.pth')) correctly calls torch.load('model.pth') inside model.load_state_dict(). Other options misuse function names or argument order.

3. Final Answer:

   model.load_state_dict(torch.load('model.pth')) -> Option A

4. Quick Check:

   Load weights with torch.load, then load_state_dict [OK]

Hint: Load file with torch.load, then pass to load_state_dict [OK]

Common Mistakes:

• Passing filename directly to load_state_dict
• Using wrong function names or order
• Confusing torch.load and load_state_dict

3. Given the code below, what will be printed?

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)

model = SimpleModel()
torch.save(model.state_dict(), 'temp.pth')

new_model = SimpleModel()
new_model.load_state_dict(torch.load('temp.pth'))

print(all(torch.equal(p1, p2) for p1, p2 in zip(model.parameters(), new_model.parameters())))

medium

A. Raises an error

B. False

C. True

D. None

Solution

1. Step 1: Understand saving and loading state_dict

   The code saves the original model's weights and loads them into a new model instance.

2. Step 2: Compare parameters of both models

   Since the new model loads the exact saved weights, parameters should be identical, so the comparison returns True.

3. Final Answer:

   True -> Option C

4. Quick Check:

   Loaded weights match saved weights = True [OK]

Hint: Loaded model matches saved weights exactly [OK]

Common Mistakes:

• Assuming new model has random weights after loading
• Thinking load_state_dict changes model architecture
• Expecting an error due to missing device argument

4. You try to load a saved state_dict into your model but get this error: RuntimeError: Error(s) in loading state_dict for Model: Missing key(s) in state_dict: "fc.weight". What is the most likely cause?

medium

A. The file path to the saved state_dict is incorrect.

B. The saved state_dict is from a different model architecture.

C. You forgot to call torch.load() before loading.

D. The model was not moved to the correct device before loading.

Solution

1. Step 1: Analyze the error message

   The error says some keys are missing in the loaded state_dict, meaning the model expects parameters not found in the saved weights.

2. Step 2: Identify cause of missing keys

   This usually happens when the saved weights come from a different model architecture than the current model.

3. Final Answer:

   The saved state_dict is from a different model architecture. -> Option B

4. Quick Check:

   Missing keys = architecture mismatch [OK]

Hint: Missing keys usually mean model architectures differ [OK]

Common Mistakes:

• Assuming file path error causes missing keys
• Forgetting to load file before loading state_dict
• Thinking device mismatch causes missing keys

5. You have a model trained on GPU and saved its state_dict. Now you want to load it on a CPU-only machine. Which code snippet correctly loads the weights without error?

hard

A. model.load_state_dict(torch.load('model_gpu.pth', device='cpu'))

B. model.load_state_dict(torch.load('model_gpu.pth'))

C. model.load_state_dict(torch.load('model_gpu.pth', map_location='cuda'))

D. model.load_state_dict(torch.load('model_gpu.pth', map_location=torch.device('cpu')))

Solution

1. Step 1: Understand device mismatch issue

   Loading GPU-trained weights on CPU requires mapping the storage to CPU to avoid errors.

2. Step 2: Use correct map_location argument

   Passing map_location=torch.device('cpu') to torch.load() correctly maps tensors to CPU.

3. Final Answer:

   model.load_state_dict(torch.load('model_gpu.pth', map_location=torch.device('cpu'))) -> Option D

4. Quick Check:

   Use map_location to load GPU weights on CPU [OK]

Hint: Use map_location=torch.device('cpu') when loading GPU weights on CPU [OK]

Common Mistakes:

• Not using map_location causes runtime errors
• Passing wrong device string like 'cuda' on CPU
• Using non-existent 'device' argument in torch.load