Solvedcloudml samples No instance key in output of canned estimator census sample

As other individuals have noted:

According to the ML Engine documentation, an instance key is required to match the returned predictions with the input data (Source)

Would it be possible to modify the canned estimator census sample to include the instance key in the output?

This was addressed on StackOverflow but it's not clear to me how to make this modification.

Thank you in advance for your help!

Ah, I see. After making this modification it works. Thanks so much for your help @elibixby.

KEY = 'gender'
def key_model_fn_gen(estimator):
    def _model_fn(features, labels, mode):
        key = features.pop(KEY)
        params = estimator.params
        model_fn_ops = estimator._model_fn(features=features, labels=labels, mode=mode, params=params)
        model_fn_ops.predictions[KEY] = key
        model_fn_ops.output_alternatives[None][1][KEY] = key # <== UPDATED
        print(model_fn_ops.output_alternatives)
        return model_fn_ops
    return _model_fn

For others' reference, this is what the model_fn_ops.output_alternatives prints out as:

 {None: (3, {'probabilities': <tf.Tensor 'binary_logistic_head/predictions/probabilities:0' shape=(?, 2) dtype=float32>, 'classes': <tf.Tensor 'binary_logistic_head/_classification_output_alternatives/classes_tensor:0' shape=(?, ?) dtype=string>, 'gender': <tf.Tensor 'ExpandDims:0' shape=(?, 1) dtype=string>})} 

generate_input_fn is for training or evaluation input functions (note that it reads data from files).
json_input_fn is for building a SavedModel binary (note that it uses placeholders).

Since you only need keys in prediction (i.e. for batch prediction or online prediction), you only need to add a key placeholder to the tensors expected by, and returned from the json_input_fn I'll add this context to the the StackOverflow answer

@elibixby Sounds like a great change to include.

@dwujellyfish - I ran into the same 'identity ... is both fed and fetched' issue in my json_serving_input_fn (but not my csv_serving_input_fn, since the CSV input records consist of just a single string Tensor field).

This is what my json_serving_input_fn looks like -- I had to pass separate 'inputs' and 'features' dicts to the tf.estimator.export.ServingInputReceiver call, with the tf.identity applied to the instance key in the 'features' dict, not the 'inputs' dict.

def json_serving_input_fn():
  inputs = {}
  features = {}

  for feat in INPUT_COLUMNS:
    inputs[feat.name] = tf.placeholder(shape=[None], dtype=feat.dtype)
    if feat.name == INSTANCE_KEY_COLUMN:
      features[feat.name] = tf.identity(inputs[feat.name])
    else:
      features[feat.name] = inputs[feat.name]

  # inputs: 'user_id': <tf.Tensor 'Placeholder:0' shape=(?,) dtype=int64>}
  # features: 'user_id': <tf.Tensor 'Identity:0' shape=(?,) dtype=int64>}

  serving_input_rcvr = tf.estimator.export.ServingInputReceiver(features, inputs)
  
  return serving_input_rcvr

Folks, I finally got around to writing a guide on how to do this. Hopefully, the two caveats at the end of the post explain the problems that people have been running into when using forward_features, and how to work around them:

https://medium.com/@lakshmanok/how-to-extend-a-canned-tensorflow-estimator-to-add-more-evaluation-metrics-and-to-pass-through-ddf66cd3047d

https://github.com/GoogleCloudPlatform/training-data-analyst/blob/master/courses/machine_learning/deepdive/07_structured/babyweight/trainer/model.py#L157