TensorFlow-gpu 2.0 - Forward features in estimators prediction output

[jusingh]

System information

• TensorFlow version (you are using): 2.0
• Are you willing to contribute it (Yes/No): No

Describe the feature and the current behavior/state.

In tensorflow-gpu 1.14, we were passing the ID key(s) as a forward feature so that it appears in prediction outputs

estimator = tf.contrib.estimator.forward_features(
estimator,
keys=[
feature_forward_key.name
for feature_forward_key in metadata.FEATURE_FORWARD_KEYS
])

On migrating the code to Tensorflow-gpu 2.0, tf.contrib package is removed and we couldn't find forward-features functionality in tensorflow core for an estimator.

Will this change the current api? How? No

Who will benefit with this feature? Everyone

Any Other info.

[karmel]

Contrib.estimator has been deprecated for quite some time. If you still need this functionality, I would recommend forking it. The code as-is is a relatively short function, and you could likely clean it up even further if it only needs to support your use-case.

[TomScheffers]

I had the same issue. The "bare-bones" solution is quite easy:

def forward_features(estimator, key):
    def new_model_fn(features, labels, mode, config):
        spec = estimator.model_fn(features, labels, mode, config)
        predictions = spec.predictions
        predictions[key] = features[key]
        spec = spec._replace(predictions=predictions)
        return spec
    return tf.estimator.Estimator(model_fn=new_model_fn, model_dir=estimator.model_dir, config=estimator.config)