Return key with tf.estimator model predictions

Question

I am using a custom tensorflow Estimator and am trying to use

tf.contrib.estimator.forward_features

to return a key column along with my outputs. I followed post1 and post2

and applied

tf.contrib.estimator.forward_features

but I am receiving the the error

 Predictions should be a dict to be able to forward features. Given: <class 'tensorflow.python.framework.ops.Tensor'>

My model functions looks like this

def model_fn(features, labels, mode):

    values = nnet(features)
    if mode == tf.estimator.ModeKeys.TRAIN:
       is_training = True
    else:
       is_training = False

    if mode == tf.estimator.ModeKeys.PREDICT:
       predictions = {
        'class_ids': tf.argmax(tf.nn.softmax(values),1),
        'probabilities': tf.nn.softmax(values),
        'logits': values,
        }
        export_outputs = {
            'prediction': tf.estimator.export.PredictOutput(predictions)
        }
    return tf.estimator.EstimatorSpec(mode,predictions=predictions,export_outputs=export_outputs)

    labels_one_hot=tf.one_hot(labels,4)

    score= tf.argmax(values,axis=1)

    loss_op= tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=values,labels=labels_one_hot))

    gradients = tf.gradients(loss_op, tf.trainable_variables())

    optimizer = tf.train.AdamOptimizer(learning_rate=.0001)

    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

    if mode == tf.estimator.ModeKeys.TRAIN:

        train_optimizer = optimizer.apply_gradients(zip(gradients, tf.trainable_variables()),

        global_step=tf.train.get_global_step())

        acc_op=tf.metrics.accuracy( labels= labels,predictions=tf.argmax(values, axis=1))

        tf.summary.scalar('accuracy_rate', acc_op[1])
        update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

        estim_specs = tf.estimator.EstimatorSpec(
             mode=mode,
             predictions=score,
             loss=loss_op,
             train_op=train_optimizer,
        eval_metric_ops={'acc': acc_op})
        return estim_specs 


    if mode == tf.estimator.ModeKeys.EVAL:
        predicted_indices = tf.argmax(values, axis=1)
        eval_metric_ops = {
          'accuracy': tf.metrics.accuracy(labels, predicted_indices)}

        return tf.estimator.EstimatorSpec(mode,loss=loss_op,eval_metric_ops=eval_metric_ops)

I am calling my estimator like this

estimator = tf.estimator.Estimator(
    model_fn= model_fn,
    config= tf.estimator.RunConfig(
            save_checkpoints_steps = 2000,
            keep_checkpoint_max = 10,
            tf_random_seed = 101),
          model_dir= "tf_dir")

estimator = tf.contrib.estimator.forward_features(
  estimator,'key')

tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

I am returning a dictionary for my predictions and my code runs fine if I don't call

tf.contrib.estimator.forward_features

score 1 · Accepted Answer · answered Aug 21 '19 at 23:25

The solution was to explicitly carry the key into model function and add it to the dictionary of outputs.

The lines of code that need to be added are

def model_fn(features, labels, mode):

    values = nnet(features)

add the key here

    key=features["key_column"]

    if mode == tf.estimator.ModeKeys.TRAIN:
        is_training = True
    else:
        is_training = False


    if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {
            'class_ids': tf.argmax(tf.nn.softmax(values),1),
             'probabilities': tf.nn.softmax(values),
             'logits': values,

add the key to the dictionary outputs

            'key_column':key
           }
        export_outputs = {
        'prediction': tf.estimator.export.PredictOutput(predictions)
          }
        return tf.estimator.EstimatorSpec(mode,predictions=predictions,export_outputs=export_outputs)

Return key with tf.estimator model predictions

1 Answers1