Saving, Loading, and Deploying PyTorch Models

# CHAPTER 16

Saving, Loading, and Deploying PyTorch Models #

1. Introduction #

2. Learning Objectives #

• Save and load model weights using state_dict.

• Save and load entire models correctly.

• Understand the importance of model.eval() during inference.

• Use TorchScript to decouple PyTorch from Python.

• Understand ONNX for cross-framework deployment.

3. Saving the State Dict (The Recommended Way) #

import torch

# Assume 'model' is a fully trained custom PyTorch model
print("Saving model weights...")

# Save only the parameters (Weights and Biases)
# .pth or .pt are the standard PyTorch file extensions
torch.save(model.state_dict(), 'my_model_weights.pth')

4. Loading the State Dict #

# 1. You MUST instantiate the exact same architecture first
loaded_model = CustomImageClassifier()

# 2. Load the dictionary of numbers from the hard drive
weights_dict = torch.load('my_model_weights.pth')

# 3. Inject the numbers into the empty architecture
loaded_model.load_state_dict(weights_dict)

# 4. CRITICAL: Set to evaluation mode before making predictions!
loaded_model.eval()

print("Model successfully loaded and ready for predictions.")

5. Saving the Entire Model (Not Recommended) #

6. Deployment Basics (REST API) #

1. 1. The web server boots up and loads your .pth file into memory.

1. 2. A user uploads a photo on the website. The website sends an HTTP POST request to your Python API.

1. 3. The API converts the image into a Tensor, calls loaded_model(tensor), and gets the result.

1. 4. The API sends the result ("It's a Dog!") back to the website via JSON.

7. TorchScript: Decoupling from Python #

# 1. Create a dummy input matching your model's expected shape
example_input = torch.rand(1, 3, 224, 224)

# 2. Trace the model (PyTorch watches how the data flows through the layers)
traced_script_module = torch.jit.trace(model, example_input)

# 3. Save the TorchScript model
traced_script_module.save("traced_model.pt")

*A C++ engineer can now load traced_model.pt directly into their C++ application using LibTorch!*

8. ONNX (Open Neural Network Exchange) #

torch.onnx.export(model,               # model being run
                  example_input,       # model input
                  "model.onnx",        # where to save the file
                  export_params=True,  # store the trained weights
                  input_names = ['input'],
                  output_names = ['output'])

9. Common Mistakes #

• Forgetting model.eval(): If you load a model and forget this line, the Dropout and BatchNorm layers will remain active. Your predictions will be random and erratic. This is the #1 bug in deployed AI models.

• Preprocessing Mismatch: If your training script scales images by dividing by 255.0, your Flask web API MUST also divide incoming user images by 255.0 before passing them to the PyTorch model.

10. Best Practices #

• Checkpointing: During long training runs, save a statedict at the end of every single Epoch (e.g., modelepoch5.pth). If your computer crashes at Epoch 49, you won't lose 48 hours of work!

1. 1. Write the code to load a saved weights file named sentimentv1.pth into a newly instantiated model named productionmodel.

1. 2. Why is saving the statedict preferred over saving the entire model?

12. MCQ Quiz with Answers #

15. Next Chapter Recommendation #