How to provide input for a TensorFlow DNNRegressor in Java?

I managed to write a TensorFlow python program with a DNNRegressor. I have trained the model and is able to get a prediction from the model in Python by manually created input (constant tensors). I have also been able to export the model in binary format.

import pandas as pd
import numpy as np
import tensorflow as tf
from tensorflow.python.framework import graph_util

#######################
# Setup
#######################

# Converting Data into Tensors
def input_fn(df, training = True):
    # Creates a dictionary mapping from each continuous feature column name (k) to
    # the values of that column stored in a constant Tensor.
    continuous_cols = {k: tf.constant(df[k].values)
                    for k in continuous_features}

    feature_cols = dict(list(continuous_cols.items()))

    if training:
        # Converts the label column into a constant Tensor.
        label = tf.constant(df[LABEL_COLUMN].values)

        # Returns the feature columns and the label.
        return feature_cols, label

    # Returns the feature columns    
    return feature_cols

def train_input_fn():
    return input_fn(train_df)

def eval_input_fn():
    return input_fn(evaluate_df)

#######################
# Data Preparation
#######################
df_train_ori = pd.read_csv('training.csv')
df_test_ori = pd.read_csv('test.csv')
train_df = df_train_ori.head(10000)
evaluate_df = df_train_ori.tail(5)
test_df = df_test_ori.head(1)
MODEL_DIR = "/tmp/model"
BIN_MODEL_DIR = "/tmp/modelBinary"
features = train_df.columns
continuous_features = [feature for feature in features if 'label' not in feature]
LABEL_COLUMN = 'label'

engineered_features = []

for continuous_feature in continuous_features:
    engineered_features.append(
        tf.contrib.layers.real_valued_column(
            column_name=continuous_feature,
            dimension=1,
            default_value=None,
            dtype=tf.int64,
            normalizer=None
            ))


#######################
# Define Our Model
#######################
regressor = tf.contrib.learn.DNNRegressor(
    feature_columns=engineered_features,
    label_dimension=1,
    hidden_units=[128, 256, 512], 
    model_dir=MODEL_DIR
    )

#######################
# Training Our Model
#######################
wrap = regressor.fit(input_fn=train_input_fn, steps=5)

#######################
# Evaluating Our Model
#######################
results = regressor.evaluate(input_fn=eval_input_fn, steps=1)
for key in sorted(results):
    print("%s: %s" % (key, results[key]))

#######################
# Save binary model (to be used in Java)
#######################
tfrecord_serving_input_fn = tf.contrib.learn.build_parsing_serving_input_fn(tf.contrib.layers.create_feature_spec_for_parsing(engineered_features)) 
regressor.export_savedmodel(
    export_dir_base=BIN_MODEL_DIR, 
    serving_input_fn = tfrecord_serving_input_fn,
    assets_extra=None,
    as_text=False,
    checkpoint_path=None,
    strip_default_attrs=False)

My next step was to load the model into java and make some predictions. I do however have a problem with specifying the input for the model in Java.

import org.tensorflow.*;
import org.tensorflow.framework.MetaGraphDef;
import org.tensorflow.framework.SignatureDef;
import org.tensorflow.framework.TensorInfo;
import java.util.List;
import java.util.Map;

public class ModelEvaluator {
    public static void main(String[] args) throws Exception {
        System.out.println("Using TF version: " + TensorFlow.version());

        SavedModelBundle model = SavedModelBundle.load("/tmp/modelBinary/1546510038", "serve");
        Session session = model.session();

        printSignature(model);
        printAllNodes(model);

        float[][] km1 = new float[1][1];
        km1[0][0] = 10;
        Tensor inKm1 = Tensor.create(km1);

        float[][] km2 = new float[1][1];
        km2[0][0] = 10000;
        Tensor inKm2 = Tensor.create(km2);

        List<Tensor<?>> outputs = session.runner()
                .feed("dnn/input_from_feature_columns/input_from_feature_columns/km1/ToFloat", inKm1)
                .feed("dnn/input_from_feature_columns/input_from_feature_columns/km2/ToFloat", inKm2)
                .fetch("dnn/regression_head/predictions/Identity:0")
                .run();

        System.out.println("\n\nOutputs from evaluation:");
        for (Tensor<?> output : outputs) {
            if (output.dataType() == DataType.STRING) {
                System.out.println(new String(output.bytesValue()));
            } else {
                float[] outArray = new float[1];
                output.copyTo(outArray);
                System.out.println(outArray[0]);
            }
        }
    }

    public static void printAllNodes(SavedModelBundle model) {
        model.graph().operations().forEachRemaining(x -> {
            System.out.println(x.name() + "   " + x.numOutputs());
        });
    }


    /**
     * This info can also be obtained from a command prompt via the command:
     * saved_model_cli show  --dir <dir-to-the-model> --tag_set serve --signature_def serving_default
     * <p>
     * See this where they also try to input data to a DNN regressor:
     * https://github.com/tensorflow/tensorflow/issues/12367
     * <p>
     * https://github.com/tensorflow/tensorflow/issues/14683
     * <p>
     * https://github.com/migueldeicaza/TensorFlowSharp/issues/293
     */
    public static void printSignature(SavedModelBundle model) throws Exception {
        MetaGraphDef m = MetaGraphDef.parseFrom(model.metaGraphDef());
        SignatureDef sig = m.getSignatureDefOrThrow("serving_default");
        int numInputs = sig.getInputsCount();
        int i = 1;
        System.out.println("-----------------------------------------------");
        System.out.println("MODEL SIGNATURE");
        System.out.println("Inputs:");
        for (Map.Entry<String, TensorInfo> entry : sig.getInputsMap().entrySet()) {
            TensorInfo t = entry.getValue();
            System.out.printf(
                    "%d of %d: %-20s (Node name in graph: %-20s, type: %s)\n",
                    i++, numInputs, entry.getKey(), t.getName(), t.getDtype());
        }
        int numOutputs = sig.getOutputsCount();
        i = 1;
        System.out.println("Outputs:");
        for (Map.Entry<String, TensorInfo> entry : sig.getOutputsMap().entrySet()) {
            TensorInfo t = entry.getValue();
            System.out.printf(
                    "%d of %d: %-20s (Node name in graph: %-20s, type: %s)\n",
                    i++, numOutputs, entry.getKey(), t.getName(), t.getDtype());
        }
        System.out.println("-----------------------------------------------");
    }
}

As can be seen from the java code I provide input for two nodes (named something with "km1" and "km2"). But I guess that is not the correct way to do it. Guess I need to provide input for the node "input_example_tensor:0"?

So question is: How do I actually create an input for the model that is loaded into java? In python I had to create a dictionary with keys "km1" and "km2", and values two constant tensors.

from How to provide input for a TensorFlow DNNRegressor in Java?