# How to Save and Load Model Checkpoints Across TensorFlow, Keras, and PyTorch

> Learn to save and load model checkpoints across TensorFlow, Keras, and PyTorch with AI-For-Beginners. Explore complete examples for seamless model persistence and reuse.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: how-to-guide
- Published: 2026-08-20

---

**The Microsoft AI‑For‑Beginners curriculum teaches checkpoint persistence through `model.save()` / `keras.models.load_model()` in TensorFlow Keras and `torch.save()` / `torch.load()` in PyTorch, with complete examples in `TransferLearningTF.ipynb` and `TransferLearningPyTorch.ipynb`.**

Saving and loading model checkpoints is essential for resuming training, deploying inference services, and versioning experiments. The **AI‑For‑Beginners** repository provides beginner‑friendly, side‑by‑side implementations in `lessons/4-ComputerVision/08-TransferLearning/`, demonstrating identical workflows for **TensorFlow Keras** and **PyTorch** using pretrained VGG16 networks fine‑tuned on a cat/dog classification task.

## TensorFlow Keras: Complete Model Serialization

In `TransferLearningTF.ipynb`, the Keras API provides a single‑line solution that persists everything—architecture, weights, and optimizer state.

### Saving a Keras Model

```python
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Build a VGG16-based transfer learning model

vgg = tf.keras.applications.VGG16(
    include_top=False,
    input_shape=(224, 224, 3),
    weights='imagenet'
)

model = keras.Sequential([
    vgg,
    layers.Flatten(),
    layers.Dense(256, activation='relu'),
    layers.Dropout(0.5),
    layers.Dense(1, activation='sigmoid')
])

model.compile(
    optimizer='adam',
    loss='binary_crossentropy',
    metrics=['accuracy']
)

# Train the model (data pipeline omitted)

# model.fit(train_ds, epochs=5, validation_data=val_ds)

# Save complete checkpoint: architecture + weights + optimizer state

model.save('data/cats_dogs.tf')  # line 859 in TransferLearningTF.ipynb

```

The `.tf` extension creates a **SavedModel** directory format. This is TF 2.x's recommended approach—portable, language‑agnostic, and compatible with TensorFlow Serving.

### Loading a Keras Model

```python

# Restore the exact same model object

restored_model = keras.models.load_model('data/cats_dogs.tf')  # line 875 in TransferLearningTF.ipynb

# Immediate inference without reconstruction

predictions = restored_model.predict(sample_batch)

```

## PyTorch: Object Serialization with torch.save()

The PyTorch workflow in `TransferLearningPyTorch.ipynb` uses Python's `pickle` protocol through `torch.save()`, serializing the entire model object.

### Saving a PyTorch Model

```python
import torch
import torchvision.models as models

# Load pretrained VGG16

vgg = models.vgg16(pretrained=True)

# Freeze feature extraction layers

for param in vgg.features.parameters():
    param.requires_grad = False

# Replace classifier head for binary classification

vgg.classifier[6] = torch.nn.Linear(4096, 1)

# Train (data loader, loss, optimizer loop omitted)

# for epoch in range(5): ...

# Save entire model object

torch.save(vgg, 'data/cats_dogs.pth')  # line 716 in TransferLearningPyTorch.ipynb

```

### Loading a PyTorch Model

```python

# Restore the complete model

vgg_restored = torch.load('data/cats_dogs.pth')

# Set to evaluation mode for inference

vgg_restored.eval()
with torch.no_grad():
    output = vgg_restored(sample_tensor)

```

## TensorFlow Keras vs. PyTorch: Checkpoint Comparison

| Aspect | TensorFlow Keras | PyTorch |
|--------|------------------|---------|
| **Save command** | `model.save('path.tf')` | `torch.save(model, 'path.pth')` |
| **Load command** | `keras.models.load_model('path.tf')` | `torch.load('path.pth')` |
| **What gets stored** | Architecture, weights, optimizer state, training config | Entire Python object (by default) |
| **File format** | SavedModel directory (protobuf + checkpoints) | Pickle‑serialized object |
| **Production best practice** | Use as‑is | Save `state_dict()` separately, reload into fresh architecture |

## Production Considerations for PyTorch Checkpointing

The **AI‑For‑Beginners** notebooks use the simplest `torch.save(model, ...)` pattern for clarity. In production code, prefer saving only the **state dictionary**:

```python

# Recommended: save just the parameters

torch.save(vgg.state_dict(), 'data/cats_dogs_weights.pth')

# Later: instantiate fresh architecture, then load weights

vgg = models.vgg16(pretrained=False)
vgg.classifier[6] = torch.nn.Linear(4096, 1)
vgg.load_state_dict(torch.load('data/cats_dogs_weights.pth'))

```

This approach:
- Decouples architecture definition from weights
- Reduces file size
- Avoids pickle compatibility issues across PyTorch versions

For cross‑device loading (training on GPU, inference on CPU), always specify `map_location`:

```python
vgg = torch.load('data/cats_dogs.pth', map_location='cpu')

```

## Summary

- **TensorFlow Keras** checkpointing uses `model.save()` and `keras.models.load_model()` in `TransferLearningTF.ipynb` to persist complete SavedModel directories containing architecture, weights, and optimizer state.

- **PyTorch** checkpointing uses `torch.save()` and `torch.load()` in `TransferLearningPyTorch.ipynb` to serialize model objects; production code should save `state_dict()` separately for better portability.

- Both approaches follow identical logical steps: build model, train, save checkpoint, load checkpoint, resume inference or training.

- Source files are located at `lessons/4-ComputerVision/08-TransferLearning/TransferLearningTF.ipynb` and `TransferLearningPyTorch.ipynb` in the microsoft/AI‑For‑Beginners repository.

## Frequently Asked Questions

### What is the difference between saving a full model and a state dict in PyTorch?

Saving the full model with `torch.save(model, ...)` pickles the entire Python object including its architecture definition. Saving `model.state_dict()` stores only the learnable parameters. The state dict approach is preferred in production because it creates smaller files, avoids version compatibility issues with pickle, and keeps architecture code explicit and version‑controlled.

### Can I load a TensorFlow Keras model without knowing its original architecture?

Yes. `keras.models.load_model()` reconstructs the complete `tf.keras.Model` from the SavedModel format without requiring access to the original Python code that defined it. This makes Keras checkpoints more portable for deployment scenarios where source code may not be available.

### How do I resume training from a checkpoint in both frameworks?

In **TensorFlow Keras**, loading restores optimizer state automatically, so you can call `model.fit()` immediately with the resumed dataset position. In **PyTorch**, `torch.load()` restores the model object but not the optimizer; you must separately save and load `optimizer.state_dict()` to resume training with the same learning rate and momentum states.

### Are these checkpoint formats compatible across different machine learning frameworks?

No. TensorFlow SavedModel and PyTorch pickle formats are not interchangeable. For framework‑agnostic model exchange, use **ONNX** (export from either framework) or convert to **TensorFlow Lite** / **TorchScript** for deployment. The **AI‑For‑Beginners** notebooks focus on native checkpointing for simplicity and educational clarity.