tflite模型调试 - 中间层output输出

原创

已于 2023-10-20 13:59:20 修改 · 3.8k 阅读

14 ·

CC 4.0 BY-SA版权

文章标签：

#tensorflow

于 2021-01-12 15:38:22 首次发布

tflite模型因内存复用，中间层output数据未保存，导致debug难点大。可通过解析flatbuffer格式的tflite模型，直接修改模型内容，调整模型output到指定层来解决。文中介绍了修改tflite模型output的函数、调用方式及修改后模型的中间层输出结果。

背景

由于tflite模型设计时考虑到了内存占用等问题，其forward的过程会反复利用中间内存空间，而导致中间层的output数据没有保存下来，模型debug难点较大。（备注：内存复用算法任务是计算整个内存区大小、中间变量重叠内存块的offset。虽然前序计算完后会被后续结果覆盖，但在实际使用中无感的。)

经过查阅相关材料，发现一个最简洁的方式，记录于此。

思路：解析flatbuffer格式的tflite模型，直接修改tflite模型内容，调整模型output到指定层。

实现

1. 修改tflite模型output的函数。

import flatbuffers
from tensorflow.lite.python import schema_py_generated as schema_fb
 
def OutputsOffset(subgraph, j):
    o = flatbuffers.number_types.UOffsetTFlags.py_type(subgraph._tab.Offset(8))
    if o != 0:
        a = subgraph._tab.Vector(o)
        return a + flatbuffers.number_types.UOffsetTFlags.py_type(j * 4)
    return 0

#参考了https://github.com/raymond-li/tflite_tensor_outputter/blob/master/tflite_tensor_outputter.py
#调整output到指定idx
def buffer_change_output_tensor_to(model_buffer, new_tensor_i):
    
    root = schema_fb.Model.GetRootAsModel(model_buffer, 0)
    output_tensor_index_offset = OutputsOffset(root.Subgraphs(0),