TFLite c++ determine the classification on output - c++

Question

I'm trying to get an output from a trained model which has a classification, the input node count is 1, and the output node count is 2. However, I'm not quite sure where the classification lands and how exactly do I handle it.
for(size_t idx = 0; idx < input_node_count; idx++)
{
float* data_ptr = interpreter->typed_input_tensor<float>(idx);
memcpy(data_ptr, my_input.data(), input_elem_size[idx]);
}
if (kTfLiteOk != interpreter->Invoke())
{
return false;
}
for(size_t idx = 0; idx < output_node_count; idx++)
{
float* output = interpreter->typed_output_tensor<float>(idx);
output_buffer[idx] = std::vector<float> (output,
output + output_elem_size[idx]);
}
result = output_buffer[1];
classification_result = output_buffer[0]; // Best way to approach this
As of now, I can just print out the sizes and see that result is 196.608 elements and classification_result is 2, as it should. My problem is I hard-coded this to be index 1 and 0 but this might not always be the case in my program which runs all sorts of models. So sometimes classification might be index 1, which causes the above code to fall apart.
I've tried to check the sizes of the buffers however that is also not guaranteed since the classification size and the result size is different for each input. Is there a way for me to know for certain which index is which? Am I approaching this the right way?

Answer 1

Use Tensorflow lite signatures for this. Signature defs can help you with that by accessing inputs/outputs using names defined in the original model.
See conversion and inference example here
Python example
# Load the TFLite model in TFLite Interpreter
interpreter = tf.lite.Interpreter(TFLITE_FILE_PATH)
# There is only 1 signature defined in the model,
# so it will return it by default.
# If there are multiple signatures then we can pass the name.
my_signature = interpreter.get_signature_runner()
# my_signature is callable with input as arguments.
output = my_signature(x=tf.constant([1.0], shape=(1,10), dtype=tf.float32))
# 'output' is dictionary with all outputs from the inference.
# In this case we have single output 'result'.
print(output['result'])
For C++
# To run
auto my_signature = interpreter_->GetSignatureRunner("my_signature");
# Set your inputs and allocate tensors
auto* input_tensor_a = my_signature->input_tensor("input_a");
...
# Execute
my_signature->Invoke();
# output
auto* output_tensor_x = my_signature->output_tensor("output_x");