Saving custom estimators in TensorFlow

Question

I am trying to save a custom estimator after training, but always receive an error. I am using TensorFlow v.1.4, and have tried various solutions I could search on the web and in tutorials and examples.

(Credit: I started following the tutorial at here , but have modified the code to suit).

Here is my code:

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Dec 16 10:17:59 2017

@author: ali
"""

import tensorflow as tf
import numpy as np
import shutil

# Define variables
SEQ_LEN = 10
DEFAULTS = [[0.0] for x in range(0, SEQ_LEN)]
BATCH_SIZE = 20
TIMESERIES_COL = 'rawdata'
N_OUTPUTS = 2  # in each sequence, 1-8 are features, and 9-10 is label
N_INPUTS = SEQ_LEN - N_OUTPUTS
N_EPOCHS = 100
LSTM_SIZE = 3  # number of hidden layers in each of the LSTM cells
LEARNING_RATE = 0.01

def create_time_series():
  freq = (np.random.random()*0.5) + 0.1  # 0.1 to 0.6
  ampl = np.random.random() + 0.5  # 0.5 to 1.5
  x = np.sin(np.arange(0,SEQ_LEN) * freq) * ampl
  return x

def to_csv(filename, N):
  with open(filename, 'w') as ofp:
    for lineno in range(0, N):
      seq = create_time_series()
      line = ",".join(map(str, seq))
      ofp.write(line + '\n')

# read data and convert to needed format
def read_dataset(filename, mode=tf.contrib.learn.ModeKeys.TRAIN):  

    def _input_fn():
        num_epochs = N_EPOCHS if mode == tf.contrib.learn.ModeKeys.TRAIN else 1

        # could be a path to one file or a file pattern.
        input_file_names = tf.train.match_filenames_once(filename)

        filename_queue = tf.train.string_input_producer(input_file_names, num_epochs=num_epochs)
        reader = tf.TextLineReader()
        _, value = reader.read_up_to(filename_queue, num_records=BATCH_SIZE)

        value_column = tf.expand_dims(value, -1)
        print('readcsv={}'.format(value_column))

        # all_data is a list of tensors
        all_data = tf.decode_csv(value_column, record_defaults=DEFAULTS)  
        inputs = all_data[:len(all_data)-N_OUTPUTS]  # first few values
        label = all_data[len(all_data)-N_OUTPUTS : ] # last few values

        # from list of tensors to tensor with one more dimension
        inputs = tf.concat(inputs, axis=1)
        label = tf.concat(label, axis=1)
        print('inputs={}'.format(inputs))

        return {TIMESERIES_COL: inputs}, label   # dict of features, label

    return _input_fn


# create the inference model
def simple_rnn(features, labels, mode, params):

    # 0. Reformat input shape to become a sequence
    x = tf.split(features[TIMESERIES_COL], N_INPUTS, 1)
    #print 'x={}'.format(x)

    # 1. configure the RNN
    lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(LSTM_SIZE, forget_bias=1.0)
    outputs, _ = tf.nn.static_rnn(lstm_cell, x, dtype=tf.float32)

    # slice to keep only the last cell of the RNN
    outputs = outputs[-1]
    #print 'last outputs={}'.format(outputs)

    # output is result of linear activation of last layer of RNN
    weight = tf.Variable(tf.random_normal([LSTM_SIZE, N_OUTPUTS]))
    bias = tf.Variable(tf.random_normal([N_OUTPUTS]))
    predictions = tf.matmul(outputs, weight) + bias

    # 2. loss function, training/eval ops
    if mode == tf.contrib.learn.ModeKeys.TRAIN or mode == tf.contrib.learn.ModeKeys.EVAL:
        loss = tf.losses.mean_squared_error(labels, predictions)
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=params["l_rate"])
        train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
        eval_metric_ops = {"rmse": tf.metrics.root_mean_squared_error(labels, predictions)}
    else:
        loss = None
        train_op = None
        eval_metric_ops = None

    # 3. Create predictions
    predictions_dict = {"predicted": predictions}

    # 4. return ModelFnOps
    return tf.estimator.EstimatorSpec(
        mode=mode,
        predictions=predictions_dict,
        loss=loss,
        train_op=train_op,
        eval_metric_ops=eval_metric_ops)


def get_train():
    return read_dataset('train.csv', mode=tf.contrib.learn.ModeKeys.TRAIN)

def get_valid():
    return read_dataset('valid.csv', mode=tf.contrib.learn.ModeKeys.EVAL)

def my_serving_input_fn():
    ''' serving input function for saving the estimator'''

    feature_spec = {TIMESERIES_COL: tf.FixedLenFeature(dtype=tf.float32, shape=[N_INPUTS])}

    serialized_tf_example = tf.placeholder(dtype=tf.string, shape=[None], name='input_example_tensor')
    receiver_tensors = {TIMESERIES_COL: serialized_tf_example}
    features = tf.parse_example(serialized_tf_example, feature_spec)

    return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
    #return tf.estimator.export.build_parsing_serving_input_receiver_fn(feature_spec)()


def generate_nn():
    model_params = {"l_rate": LEARNING_RATE}
    nn = tf.estimator.Estimator(model_fn=simple_rnn, params=model_params, model_dir='./output_dir')
    return nn


def save_nn(nn_estimator, output_dir):
    nn_estimator.export_savedmodel(output_dir, my_serving_input_fn)
    print('Successfully saved the estimator...')


def main():

    # remove previous files
    shutil.rmtree('output_dir', ignore_errors=True) 
    shutil.rmtree('test_dir', ignore_errors=True)

    # generate data
    to_csv('train.csv', 5000)
    to_csv('test.csv', 1000)

    # instantiate the nn estimator
    nn = generate_nn()

    # train nn
    nn.train(get_train(), steps=2000)

    # evaluate nn
    ev = nn.evaluate(input_fn=get_valid())
    print(ev)

    # save nn for future use
    save_nn(nn, './test_dir')

if __name__ == '__main__':
    main()

Here is the error I receive:

   File "/.../RNN-estimators-v3.py", line 172, in <module>
main()

  File "/.../RNN-estimators-v3.py", line 167, in main
save_nn(nn, './test_dir')

  File "/.../RNN-estimators-v3.py", line 142, in save_nn
nn_estimator.export_savedmodel(output_dir, my_serving_input_fn)

  File "/.../anaconda/envs/TF-1-4-CPU/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 534, in export_savedmodel
serving_input_receiver.receiver_tensors_alternatives)

  File "/.../anaconda/envs/TF-1-4-CPU/lib/python3.6/site-packages/tensorflow/python/estimator/export/export.py", line 195, in build_all_signature_defs
'{}'.format(type(export_outputs)))

ValueError: export_outputs must be a dict and not<class 'NoneType'>

Your help is much appreciated.

Answer 1

Make sure to include export_outputs in your model_fn function when the mode is Predict.

def simple_rnn(features, labels, mode, params):

    # 0. Reformat input shape to become a sequence
    x = tf.split(features[TIMESERIES_COL], N_INPUTS, 1)
    #print 'x={}'.format(x)

    # 1. configure the RNN
    lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(LSTM_SIZE, forget_bias=1.0)
    outputs, _ = tf.nn.static_rnn(lstm_cell, x, dtype=tf.float32)

    # slice to keep only the last cell of the RNN
    outputs = outputs[-1]
    #print 'last outputs={}'.format(outputs)

    # output is result of linear activation of last layer of RNN
    weight = tf.Variable(tf.random_normal([LSTM_SIZE, N_OUTPUTS]))
    bias = tf.Variable(tf.random_normal([N_OUTPUTS]))
    predictions = tf.matmul(outputs, weight) + bias

    # 2. loss function, training/eval ops
    if mode == tf.contrib.learn.ModeKeys.TRAIN or mode == tf.contrib.learn.ModeKeys.EVAL:
        loss = tf.losses.mean_squared_error(labels, predictions)
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=params["l_rate"])
        train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
        eval_metric_ops = {"rmse": tf.metrics.root_mean_squared_error(labels, predictions)}
        return  tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            train_op=train_op,
            eval_metric_ops=eval_metric_ops)

    else:
        loss = None
        train_op = None
        eval_metric_ops = None

        # 3. Create predictions

        export_outputs = {'predict_output': tf.estimator.export.PredictOutput({"pred_output_classes": predictions, 'probabilities': #your probabilities})}
        predictions_dict = {"predicted": predictions}
        # 4. return ModelFnOps

        return tf.estimator.EstimatorSpec(
            mode=mode,
            predictions=predictions_dict,
            loss=loss,
            train_op=train_op,
            eval_metric_ops=eval_metric_ops,export_outputs=export_outputs )

Answer 2

When exporting a graph, there is a required export_outputs field in the EstimatorSpec . See the model_fn documentation for details.

I'd also note that tf.contrib.timeseries has some of this boilerplate written for you (including an RNN example ).