Saving Model Parameters

During training, it is esssential to persist learned parameters to disk for later validation, inference, or continued training. TensorFlow provides the tf.train.Saver class for this purpose.

To begin, instantiate a Saver object:

saver = tf.train.Saver()

The max_to_keep parameter controls how many checkpoint files are retained. By default, the five most recent checkpoints are kept:

saver = tf.train.Saver(max_to_keep=5)

Setting max_to_keep to None or 0 saves a checkpoint at every step, though this consumes significant disk space and is rarely necessary:

saver = tf.train.Saver(max_to_keep=0)

To preserve only the most recent checkpoint, set max_to_keep to 1:

saver = tf.train.Saver(max_to_keep=1)

Call the save method to persist the current session state:

saver.save(sess, 'checkpoints/model', global_step=epoch)

The second argument specfiies the file path and base name. The global_step suffix gets appended automatically:

• saver.save(sess, 'my-model', global_step=0) → 'my-model-0'
• saver.save(sess, 'my-model', global_step=1000) → 'my-model-1000'

Complete Example

import tensorflow as tf
import numpy as np

# Define computation graph
input_tensor = tf.placeholder(tf.float32, shape=[None, 1])
ground_truth = 5 * input_tensor + 2

# Model parameters
weights = tf.Variable(tf.random_normal([1], -1, 1))
biases = tf.Variable(tf.zeros([1]))
prediction = weights * input_tensor + biases

# Loss and optimizer
mse = tf.reduce_mean(tf.square(ground_truth - prediction))
train_op = tf.train.GradientDescentOptimizer(0.3).minimize(mse)

# Training configuration
training_mode = False
total_iterations = 200
save_interval = 40
checkpoint_dir = './model_ckpts/'

# Create saver
saver = tf.train.Saver()

# Generate sample data
sample_data = np.reshape(np.random.rand(15).astype(np.float32), (15, 1))

with tf.Session() as session:
    session.run(tf.initialize_all_variables())
    
    if training_mode:
        for iteration in range(total_iterations):
            session.run(train_op, feed_dict={input_tensor: sample_data})
            if (iteration + 1) % save_interval == 0:
                saver.save(session, checkpoint_dir + 'linear_model.ckpt', 
                          global_step=iteration + 1)
    else:
        # Restore from checkpoint
        state = tf.train.get_checkpoint_state(checkpoint_dir)
        if state and state.model_checkpoint_path:
            saver.restore(session, state.model_checkpoint_path)
        else:
            print("No checkpoint found")
        
        print("Weights:", session.run(weights))
        print("Biases:", session.run(biases))

Restoring Model State

Use the restore method to reload previously saved variables:

saver.restore(sess, 'checkpoints/model-100')

The first argument is the target session where parameters will be loaded. The second argument points to the checkpoint file path. Unlike the save operation, you do not need to specify the checkpoint filename—the Saver reads the checkpoint metadata file to locate the most recent valid checkpoint.