Saving model state_dict in PyTorch - Model Metrics & Evaluation

Metrics & Evaluation - Saving model state_dict

Which metric matters for this concept and WHY

When saving a model's state_dict in PyTorch, the key metric is model reproducibility. This means you can reload the saved weights exactly and get the same predictions. The saved state_dict contains all learned parameters (weights and biases) of the model. Ensuring it is saved and loaded correctly guarantees consistent model performance.

Confusion matrix or equivalent visualization (ASCII)

Saving state_dict is not about classification metrics, but about preserving model parameters. However, to check if saving/loading worked, you can compare predictions before and after saving:

    Before saving:  [0, 1, 1, 0, 1]
    After loading:  [0, 1, 1, 0, 1]
    Match: True
    

If predictions match exactly, the state_dict saved and loaded correctly.

Precision vs Recall (or equivalent tradeoff) with concrete examples

Saving state_dict is about exactness, not tradeoffs like precision or recall. But consider this analogy:

• Saving too little: If you save only part of the state_dict, the model will lose information, like low recall (missing important parts).
• Saving too much: Saving extra unnecessary data can make files large but doesn't harm accuracy, like high precision but low recall.

Best practice is to save the complete state_dict for full model recovery.

What "good" vs "bad" metric values look like for this use case

Good outcome:

• Model predictions before saving and after loading match exactly.
• File size is reasonable, containing only model parameters.
• No errors when loading the state_dict.

Bad outcome:

• Predictions differ after loading, indicating corrupted or incomplete save.
• File is too large or missing parameters.
• Loading throws errors or mismatches model architecture.

Metrics pitfalls (accuracy paradox, data leakage, overfitting indicators)

• Saving incomplete state_dict: Forgetting to save optimizer state or parts of the model can cause training to fail on reload.
• Architecture mismatch: Loading a state_dict into a different model structure causes errors.
• Overwriting files: Accidentally overwriting good saved models with bad ones loses progress.
• Data leakage: Not related here, but ensure saved model is tested on unseen data after loading.

Self-check

Your model has 98% accuracy before saving. After loading the state_dict, predictions drop to 70%. Is it good?

Answer: No, this means the state_dict was not saved or loaded correctly. The model parameters changed or were corrupted. You should verify saving/loading code and ensure the model architecture matches exactly.

Key Result

Saving and loading the complete model state_dict ensures exact reproducibility of model predictions.

Practice

1. What does model.state_dict() in PyTorch contain?

easy

A. Only the optimizer settings

B. The learned parameters (weights and biases) of the model

C. The entire model architecture and code

D. The training dataset

Solution

1. Step 1: Understand what state_dict holds

   The state_dict stores all the learned parameters like weights and biases of the model layers.

2. Step 2: Differentiate from other components

   It does not include the model architecture code or optimizer settings, only the parameters.

3. Final Answer:

   The learned parameters (weights and biases) of the model -> Option B

4. Quick Check:

   state_dict = learned parameters [OK]

Hint: state_dict always means model weights only [OK]

Common Mistakes:

• Thinking state_dict saves the whole model code
• Confusing optimizer state with model state
• Assuming it saves the training data

2. Which of the following is the correct syntax to save a PyTorch model's state_dict to a file named 'model.pth'?

easy

A. torch.save(model.state_dict(), 'model.pth')

B. model.state_dict().save('model.pth')

C. model.save_state('model.pth')

D. torch.save(model, 'model.pth')

Solution

1. Step 1: Recall the saving function

   In PyTorch, torch.save() is used to save objects to a file.

2. Step 2: Save only the state_dict

   To save the model parameters, you pass model.state_dict() to torch.save() along with the filename.

3. Final Answer:

   torch.save(model.state_dict(), 'model.pth') -> Option A

4. Quick Check:

   Use torch.save with state_dict [OK]

Hint: Use torch.save(model.state_dict(), filename) to save weights [OK]

Common Mistakes:

• Saving the whole model instead of state_dict
• Using non-existent save_state method
• Calling save on state_dict directly

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(), 'weights.pth')
loaded_state = torch.load('weights.pth')
print(type(loaded_state))

medium

A.

B.

C.

D.

Solution

1. Step 1: Understand what torch.save stores

   Saving model.state_dict() stores an OrderedDict of parameter tensors.

2. Step 2: Loading with torch.load returns the same type

   When loaded, it returns an OrderedDict, not a Module or plain dict.

3. Final Answer:

   <class 'collections.OrderedDict'> -> Option C

4. Quick Check:

   state_dict loads as OrderedDict [OK]

Hint: state_dict loads as OrderedDict, not model or tensor [OK]

Common Mistakes:

• Expecting loaded_state to be a model instance
• Thinking it returns a plain dict
• Confusing with tensor type

4. You saved a model's state_dict with torch.save(model.state_dict(), 'model.pth'). Later, you try to load it with model.load_state_dict(torch.load('model.pth')) but get a runtime error about missing keys. What is the most likely cause?

medium

A. The model architecture does not match the saved state_dict

B. The file 'model.pth' is corrupted

C. You forgot to call model.eval() before loading

D. You used torch.save(model, 'model.pth') instead

Solution

1. Step 1: Understand load_state_dict requirements

   Loading weights requires the model architecture to match the saved parameters exactly.

2. Step 2: Identify cause of missing keys error

   If keys are missing, it usually means the model layers differ from those saved in the state_dict.

3. Final Answer:

   The model architecture does not match the saved state_dict -> Option A

4. Quick Check:

   Mismatch architecture causes missing keys error [OK]

Hint: Check model matches saved weights architecture [OK]

Common Mistakes:

• Assuming file corruption without checking
• Thinking eval mode affects loading
• Confusing saving whole model vs state_dict

5. You want to save a PyTorch model's state_dict after training and later load it to continue training on a different machine. Which sequence of steps is correct?

hard

A. Save model.state_dict() and optimizer.state_dict() together in one file, then load both on new machine

B. Save with torch.save(model, 'file.pth'), then load with model = torch.load('file.pth') without defining architecture

C. Save optimizer state only, then recreate model and optimizer on new machine

D. Save with torch.save(model.state_dict(), 'file.pth'), then on new machine create same model architecture and load with model.load_state_dict(torch.load('file.pth'))

Solution

1. Step 1: Save only model parameters

   Use torch.save(model.state_dict(), 'file.pth') to save learned weights.

2. Step 2: Recreate model architecture on new machine

   Define the same model class and create an instance before loading weights.

3. Step 3: Load saved weights into model

   Use model.load_state_dict(torch.load('file.pth')) to load parameters.

4. Final Answer:

   Save state_dict, recreate model, then load state_dict -> Option D

5. Quick Check:

   Save weights, recreate model, load weights [OK]

Hint: Always recreate model before loading state_dict [OK]

Common Mistakes:

• Trying to load weights without model definition
• Saving whole model causing compatibility issues
• Ignoring optimizer state when continuing training