How to pass key to Tensorflow graph

When you are using Google Cloud AI Platform batch prediction, launching a job takes an input path (where multiple inputs are placed) and an output path (where it will put the result in files).

The problem is that you can’t link the output result to its input, so I needed to pass the key in input to the model and each result will have the key allowing to find back which input it was.

I couldn't find a solution on the internet myself so I had to implement it modifying Tensorflow Models tools, I hope it will help people that was in the same case than me.

Tensorflow Models tool to export models

My modified version of their tool to add a key to both input and output

Differences with original

How to use it:

Initial model graph

Modified graph, notice the input_keys and output_keys

Additionally, I have built a notebook to download, edit and deploy a model to AI Platform with Google Colab

UPDATE: i’ve built a Cloud Run component that does the work here