Saving entire model in PyTorch - ML Experiment: Train & Evaluate

Experiment - Saving entire model

Problem:You have trained a PyTorch model but only saved the model's state dictionary. When loading, you face issues because the model architecture code must be available and consistent.

Current Metrics:Model trains with 90% accuracy on training data and 85% on validation data. Model saving uses torch.save(model.state_dict(), 'model.pth').

Issue:Saving only the state dictionary requires re-defining the model architecture when loading, which can cause errors or inconsistencies if the code changes.

Your Task

Save the entire PyTorch model including architecture and weights so it can be loaded directly.

Use PyTorch's recommended methods for saving and loading entire models.

Do not change the model architecture or training code.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

# Define a simple model
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc = nn.Linear(10, 2)
    def forward(self, x):
        return self.fc(x)

# Create model instance
model = SimpleNet()

# Dummy training loop (skipped for brevity)
# Assume model is trained here

# Save entire model
torch.save(model.cpu(), 'model_full.pth')

# Later or in another script: Load entire model
loaded_model = torch.load('model_full.pth')
loaded_model.eval()

# Test loaded model with dummy input
input_tensor = torch.randn(1, 10)
output = loaded_model(input_tensor)
print('Model output:', output)

Changed saving method from saving only state_dict to saving entire model with torch.save(model, path).

Used torch.load to load the entire model.

Moved model to CPU before saving to ensure portability.

Results Interpretation

Before: Saved only state_dict. Needed model class code to load. Risk of mismatch or errors.

After: Saved entire model. Can load directly. Easier and safer for deployment.

Saving the entire model in PyTorch allows you to reload it easily, reducing errors and simplifying deployment.

Bonus Experiment

Try saving and loading the model's state_dict only, then re-define the model class and load the weights. Compare this approach with saving the entire model.

💡 Hint

Use model.state_dict() to save and model.load_state_dict() to load weights. Ensure the model architecture matches exactly.