Loading model state_dict in PyTorch - ML Experiment: Train & Evaluate

Experiment - Loading model state_dict

Problem:You have trained a PyTorch model and saved its state_dict. Now you want to load this saved state_dict into a new model instance to continue training or make predictions.

Current Metrics:N/A (focus is on correct loading of model weights)

Issue:If the state_dict is not loaded correctly, the model will have random weights, leading to poor predictions or errors.

Your Task

Load the saved state_dict into a new model instance correctly and verify that the model produces the same output as the original model for the same input.

Do not retrain the model.

Use the same model architecture for loading the state_dict.

Use PyTorch's recommended methods for saving and loading.

Hint 1

Hint 2

Hint 3

Solution

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(5, 2)
    def forward(self, x):
        return self.linear(x)

# Create and train a model (simulate training by setting weights manually for reproducibility)
original_model = SimpleModel()
with torch.no_grad():
    original_model.linear.weight.fill_(1.0)
    original_model.linear.bias.fill_(0.5)

# Save the state_dict
torch.save(original_model.state_dict(), 'model_state.pth')

# Create a new model instance
loaded_model = SimpleModel()

# Load the saved state_dict
state_dict = torch.load('model_state.pth')
loaded_model.load_state_dict(state_dict)

# Set both models to eval mode
original_model.eval()
loaded_model.eval()

# Test input
input_tensor = torch.ones(1, 5)

# Get outputs
original_output = original_model(input_tensor)
loaded_output = loaded_model(input_tensor)

# Check if outputs are the same
outputs_match = torch.allclose(original_output, loaded_output)

print(f"Outputs match: {outputs_match}")
print(f"Original output: {original_output}")
print(f"Loaded output: {loaded_output}")

Defined a simple model architecture.

Simulated training by manually setting weights and bias for reproducibility.

Saved the model's state_dict to a file.

Created a new model instance with the same architecture.

Loaded the saved state_dict into the new model.

Verified that outputs from original and loaded models match for the same input.

Results Interpretation

Before loading the state_dict, the new model has random weights and produces different outputs from the original model.

After loading the state_dict, the new model produces the same outputs as the original model for the same input.

Loading a saved state_dict correctly restores the model's learned weights, allowing the model to produce consistent outputs without retraining.

Bonus Experiment

Try saving and loading the entire model using torch.save(model) and torch.load(), then compare outputs with the state_dict method.

💡 Hint

Saving the entire model saves the architecture and weights but is less flexible and less recommended than saving 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