使用实例密钥进行训练和预测

我能够训练我的模型并使用ML Engine进行预测，但是我的结果不包含任何识别信息。一次提交一行进行预测时，这很好用，但是当提交多行时，我无法将预测连接回原始输入数据。该GCP文件讨论了使用实例密钥，但我找不到任何示例代码火车和预测使用实例密钥。以GCP人口普查为例，我将如何更新输入函数以通过图形传递唯一ID，并在训练过程中忽略它，同时返回带有预测的唯一ID？或者，如果有人知道其他示例已经在使用键也将有所帮助。

来自人口普查估算器样本

def serving_input_fn():
    feature_placeholders = {
      column.name: tf.placeholder(column.dtype, [None])
      for column in INPUT_COLUMNS
    }

    features = {
      key: tf.expand_dims(tensor, -1)
      for key, tensor in feature_placeholders.items()
    }

    return input_fn_utils.InputFnOps(
      features,
      None,
      feature_placeholders
    )


def generate_input_fn(filenames,
                  num_epochs=None,
                  shuffle=True,
                  skip_header_lines=0,
                  batch_size=40):

    def _input_fn():
        files = tf.concat([
          tf.train.match_filenames_once(filename)
          for filename in filenames
        ], axis=0)

        filename_queue = tf.train.string_input_producer(
          files, num_epochs=num_epochs, shuffle=shuffle)
        reader = tf.TextLineReader(skip_header_lines=skip_header_lines)

        _, rows = reader.read_up_to(filename_queue, num_records=batch_size)

        row_columns = tf.expand_dims(rows, -1)
        columns = tf.decode_csv(row_columns, record_defaults=CSV_COLUMN_DEFAULTS)
        features = dict(zip(CSV_COLUMNS, columns))

        # Remove unused columns
        for col in UNUSED_COLUMNS:
          features.pop(col)

        if shuffle:
           features = tf.train.shuffle_batch(
             features,
             batch_size,
             capacity=batch_size * 10,
             min_after_dequeue=batch_size*2 + 1,
             num_threads=multiprocessing.cpu_count(),
             enqueue_many=True,
             allow_smaller_final_batch=True
           )
        label_tensor = parse_label_column(features.pop(LABEL_COLUMN))
        return features, label_tensor

    return _input_fn

更新：我能够使用下面这个答案中的建议代码，我只需要稍加更改即可更新model_fn_ops中的输出替代品，而不仅仅是预测字典。但是，仅当我的服务输入函数被编码为类似于this的json输入时，此方法才有效。我的服务输入功能先前是根据Census Core Sample中的CSV服务输入功能建模的。

我认为我的问题来自build_standardized_signature_def函数，甚至更多来自它调用的is_classification_problem函数。使用csv服务功能的输入字典长度为1，因此此逻辑最终使用category_signature_def结束，该分类结果仅显示分数（结果实际上是概率），而使用json服务输入功能时，输入字典长度大于1而是使用了包含所有输出的predict_signature_def。