Save/load models

Saving and loading models is a crucial step in deploying deep learning applications. It allows you to persist a trained model’s architecture and learned parameters to disk so you can reuse it later for inference or further training without retraining from scratch.

Different frameworks provide various methods for model serialization. For example, TensorFlow and Keras use the save() and load_model() functions, while PyTorch uses torch.save() and torch.load() for saving model state dictionaries or entire models.

Example: Saving and loading models in TensorFlow/Keras

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

# Define a simple model
model = Sequential([
    Dense(64, activation='relu', input_shape=(100,)),
    Dense(10, activation='softmax')
])

# Compile and train the model
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
# model.fit(x_train, y_train, epochs=10)  # Assume training happens here

# Save the entire model
model.save('my_keras_model.h5')

# Load the saved model
loaded_model = tf.keras.models.load_model('my_keras_model.h5')
        

Example: Saving and loading models in PyTorch

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(100, 10)

    def forward(self, x):
        return self.fc(x)

model = SimpleModel()

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

# Load model state_dict into a new model instance
loaded_model = SimpleModel()
loaded_model.load_state_dict(torch.load('model_state.pth'))
loaded_model.eval()