Example: Save and Load a TensorFlow Model

2020 update: just a heads up, this post is about TensorFlow 1.x. For TensorFlow 2.x, you probably want to check out this guide .

This post details how to save and load a TensorFlow model using the DNNClassifier API.

The key idea here is that you define a function or a class beforehand that takes a model directory (in which it will save and restore the model parameters), adds that to RunConfig, and returns a tf.contrib.learn.Estimator, for example, tf.contrib.learn.DNNClassifier. See make_estimator for more details.

import numpy as np
import tensorflow as tf

from tensorflow.contrib.learn.python.learn.estimators import run_config
from tensorflow.contrib.training.python.training import hparam
MODEL_DIR = 'your-model-dir'
hparams = hparam.HParams(
    num_epochs=10, 
    train_batch_size=50, 
    eval_batch_size=50,
    eval_steps=10
)
def make_estimator(model_dir):
    config = run_config.RunConfig(model_dir=model_dir)
    
    input_columns = [
      tf.feature_column.numeric_column(key='random_feature'),
    ]

    return tf.contrib.learn.DNNClassifier(
        config=config,
        n_classes=2,
        feature_columns=input_columns,
        hidden_units=[1024, 512, 256],
    )
dataset_size = 1000
X_train = np.random.rand(dataset_size)
y_train = np.random.rand(dataset_size)
X_eval = np.random.rand(dataset_size)
y_eval = np.random.rand(dataset_size)

estimator = make_estimator(MODEL_DIR)

experiment = tf.contrib.learn.Experiment(
    estimator,
    train_input_fn=tf.estimator.inputs.numpy_input_fn(
        x={'random_feature': X_train},
        y=y_train,
        num_epochs=hparams.num_epochs,
        batch_size=hparams.train_batch_size,
        shuffle=True
    ),
    eval_input_fn=tf.estimator.inputs.numpy_input_fn(
        x={'random_feature': X_eval},
        y=y_eval,
        num_epochs=None,
        batch_size=hparams.eval_batch_size,
        shuffle=False # Don't shuffle evaluation data
    )
)
experiment.train()

Loading the model

estimator_from_file = make_estimator(MODEL_DIR)
X_predict = np.array([0.3, 0.4, 0.5])

predict_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={'random_feature': X_predict},
    num_epochs=1,
    batch_size=hparams.eval_batch_size,
    shuffle=False # Don't shuffle evaluation data
)

predictions = estimator_from_file.predict_proba(input_fn=predict_input_fn)
for prediction in predictions:
    print('Prediction:', prediction)

Contents (top)