Visualize TFLite graph and get intermediate values of a particular node?

Question

I was wondering if there is a way to know the list of inputs and outputs for a particular node in tflite? I know that I can get input/outputs details, but this does not allow me to reconstruct the computation process that happens inside an Interpreter. So what I do is:

interpreter = tf.lite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.get_tensor_details()

The last 3 commands basically give me dictionaries which don't seem to have the necessary information.

So I was wondering if there is way to know where each nodes outputs goes? Surely Interpreter knows this somehow. Can we? Thanks.

Answer 1

Note: this answer was written for Tensorflow 1.x and, while the concept and core idea remains the same in TensorFlow 2.x, the commands in this answer might be deprecated.

The mechanism of TF-Lite makes the whole process of inspecting the graph and getting the intermediate values of inner nodes a bit tricky. The get_tensor(...) method suggested by the other answer does not work.

How to visualize TF-Lite inference graph?

TensorFlow Lite models can be visualized using the visualize.py script in the TensorFlow Lite repository. You just need to:

• Clone the TensorFlow repository

• Run the visualize.py script with bazel:

    bazel run //tensorflow/lite/tools:visualize \
         model.tflite \
         visualized_model.html
  

Does the nodes in my TF model have a equivalent one in TF-Lite?

NO! In fact, TF-Lite can modify your graph so that it become more optimal. Here are some words about it from the TF-Lite documentation:

A number of TensorFlow operations can be processed by TensorFlow Lite even though they have no direct equivalent. This is the case for operations that can be simply removed from the graph (tf.identity), replaced by tensors (tf.placeholder), or fused into more complex operations (tf.nn.bias_add). Even some supported operations may sometimes be removed through one of these processes.

Moreover, the TF-Lite API currently doesn't allow to get node correspondence; it's hard to interpret the inner format of TF-Lite. So, you can't get the intermediate outputs for any nodes you want, even without the one more issue below...

Can I get intermediate values of some TF-Lite nodes?

NO! Here, I will explain why get_tensor(...) wouldn't work in TF-Lite. Suppose in the inner representation, the graph contains of 3 tensors, together with some dense operations (nodes) in-between (you can think of tensor1 as input and tensor3 as output of your model). During inference of this particular graph, TF-Lite only needs 2 buffers, let's show how.

First, use tensor1 to compute tensor2 by applying dense operation. This only requires 2 buffers to store the values:

           dense              dense
[tensor1] -------> [tensor2] -------> [tensor3]
 ^^^^^^^            ^^^^^^^
 bufferA            bufferB

Second, use the value of tensor2 stored in bufferB to compute tensor3... but wait! We don't need bufferA anymore, so let's use it to store the value of tensor3:

           dense              dense
[tensor1] -------> [tensor2] -------> [tensor3]
                    ^^^^^^^            ^^^^^^^
                    bufferB            bufferA

Now is the tricky part. The "output value" of tensor1 will still point to bufferA, which now holds the values of tensor3. So if you call get_tensor(...) for the 1st tensor, you'll get incorrect values. The documentation of this method even states:

This function cannot be used to read intermediate results.

How to get around this?

• Easy but limited way. You can specify the names of the nodes, output tensors of which you want to get the values of during conversion:

    tflite_convert \
        -- # other options of your model
        --output_arrays="output_node,intermediate/node/n1,intermediate/node/n2"
  

• Hard but flexible way. You can compile TF-Lite with Bazel (using this instruction). Then you can actually inject some logging code to Interpreter::Invoke() in the file tensorflow/lite/interpreter.cc. An ugly hack, but it works.

Answer 2

As @FalconUA has pointed out, we cannot directly get intermediate inputs and outputs from a TFlite model. But, we can get inputs and outputs of layers by modifying the model buffer. This repo shows how it is done. We need to modify flat buffer schema for this to work. The modified TFlite schema (tflite folder in the repo) is available in the repo.

For the completeness of the answer, below is the relevant code:

def buffer_change_output_tensor_to(model_buffer, new_tensor_i):
    # from https://github.com/raymond-li/tflite_tensor_outputter
    # Set subgraph 0's output(s) to new_tensor_i
    # Reads model_buffer as a proper flatbuffer file and gets the offset programatically
    # It might be much more efficient if Model.subgraphs[0].outputs[] was set to a list of all the tensor indices.
    fb_model_root = tflite_model.Model.GetRootAsModel(model_buffer, 0)
    output_tensor_index_offset = fb_model_root.Subgraphs(0).OutputsOffset(0) # Custom added function to return the file offset to this vector
    # print("buffer_change_output_tensor_to. output_tensor_index_offset: ")
    # print(output_tensor_index_offset)
    # output_tensor_index_offset = 0x5ae07e0 # address offset specific to inception_v3.tflite
    # output_tensor_index_offset = 0x16C5A5c # address offset specific to inception_v3_quant.tflite
    # Flatbuffer scalars are stored in little-endian.
    new_tensor_i_bytes = bytes([
        new_tensor_i & 0x000000FF, \
        (new_tensor_i & 0x0000FF00) >> 8, \
        (new_tensor_i & 0x00FF0000) >> 16, \
        (new_tensor_i & 0xFF000000) >> 24 \
    ])
    # Replace the 4 bytes corresponding to the first output tensor index
    return model_buffer[:output_tensor_index_offset] + new_tensor_i_bytes + model_buffer[output_tensor_index_offset + 4:]

def get_tensor(path_tflite, tensor_id):
    with open(path_tflite, 'rb') as fp:
        model_buffer = fp.read()
    
    model_buffer = buffer_change_output_tensor_to(model_buffer, int(tensor_id))
    interpreter = tf.lite.Interpreter(model_content=model_buffer)
    interpreter.allocate_tensors()
    tensor_details = interpreter._get_tensor_details(tensor_id)
    tensor_name = tensor_details['name']
    
    input_details = interpreter.get_input_details()
    interpreter.set_tensor(input_details[0]['index'], input_tensor)
    interpreter.invoke()
    
    tensor = interpreter.get_tensor(tensor_id)
    return tensor