#!/usr/bin/env python3 """dump TFLite 模型前 N 个 operator 的结构与常量(perm/reshape shape)""" import struct import sys TENSOR_TYPE = {0: "FLOAT32", 1: "FLOAT16", 2: "INT32", 3: "UINT8", 4: "INT64", 5: "STRING", 6: "BOOL", 7: "INT16", 8: "COMPLEX64", 9: "INT8", 10: "FLOAT64", 11: "COMPLEX128", 12: "UINT64", 13: "RESOURCE", 14: "VARIANT", 15: "UINT32", 16: "UINT16", 17: "INT4", 18: "BFLOAT16", 19: "FLOAT8_E4M3FN", 20: "FLOAT8_E4M3FNUZ", 21: "FLOAT8_E5M2", 22: "FLOAT8_E5M2FNUZ"} BUILTIN = {0: "ADD", 1: "AVERAGE_POOL_2D", 2: "CONCATENATION", 3: "CONV_2D", 4: "DEPTHWISE_CONV_2D", 5: "DEPTH_TO_SPACE", 6: "DEQUANTIZE", 7: "EMBEDDING_LOOKUP", 8: "FLOOR", 9: "FULLY_CONNECTED", 10: "HASHTABLE_LOOKUP", 11: "L2_NORMALIZATION", 12: "L2_POOL_2D", 13: "LOCAL_RESPONSE_NORMALIZATION", 14: "LOGISTIC", 15: "LSH_PROJECTION", 16: "LSTM", 17: "MAX_POOL_2D", 18: "MUL", 19: "RELU", 20: "RELU_N1_TO_1", 21: "RELU6", 22: "RESHAPE", 23: "RESIZE_BILINEAR", 24: "RNN", 25: "SOFTMAX", 26: "SPACE_TO_DEPTH", 27: "SVDF", 28: "TANH", 29: "CONCAT_EMBEDDINGS", 30: "SKIP_GRAM", 31: "CALL", 32: "CUSTOM", 33: "EMBEDDING_LOOKUP_SPARSE", 34: "PAD", 35: "UNIDIRECTIONAL_SEQUENCE_RNN", 36: "GATHER", 37: "BATCH_TO_SPACE_ND", 38: "SPACE_TO_BATCH_ND", 39: "TRANSPOSE", 40: "MEAN", 41: "SUB", 42: "DIV", 43: "SQUEEZE", 44: "UNIDIRECTIONAL_SEQUENCE_LSTM", 45: "STRIDED_SLICE", 46: "BIDIRECTIONAL_SEQUENCE_RNN", 47: "EXP", 48: "TOPK_V2", 49: "SPLIT", 50: "LOG_SOFTMAX", 51: "DELEGATE", 52: "BIDIRECTIONAL_SEQUENCE_LSTM", 53: "CAST", 54: "PREDICT", 55: "CONCATENATION_RELU", 56: "ARG_MAX", 57: "MINIMUM", 58: "LESS", 59: "NEG", 60: "PADV2", 61: "GREATER", 62: "GREATER_EQUAL", 63: "LESS_EQUAL", 64: "SELECT", 65: "SLICE", 66: "SIN", 67: "TRANSPOSE_CONV", 68: "SPARSE_TO_DENSE", 69: "TILE", 70: "EXPAND_DIMS", 71: "EQUAL", 72: "NOT_EQUAL", 73: "LOG", 74: "SUM", 75: "SQRT", 76: "RSQRT", 77: "SHAPE", 78: "POW", 79: "ARG_MIN", 80: "FAKE_QUANT", 81: "REDUCE_PROD", 82: "REDUCE_MAX", 83: "PACK", 84: "LOGICAL_OR", 85: "LOGICAL_AND", 86: "LOGICAL_NOT", 87: "UNPACK", 88: "REDUCE_MIN", 89: "FLOOR_DIV", 90: "REDUCE_ANY", 91: "SQUARE", 92: "ZEROS_LIKE", 93: "FILL", 94: "FLOOR_MOD", 95: "RANGE", 96: "RESIZE_NEAREST_NEIGHBOR", 97: "LEAKY_RELU", 98: "SQUARED_DIFFERENCE", 99: "MIRROR_PAD", 100: "ABS", 101: "SPLIT_V", 102: "UNIQUE", 103: "CEIL", 104: "REVERSE_V2", 105: "ADD_N", 106: "GATHER_ND", 107: "COS", 108: "WHERE", 109: "RANK", 110: "ELU", 111: "REVERSE_SEQUENCE", 112: "MATRIX_DIAG", 113: "QUANTIZE", 114: "MATRIX_SET_DIAG", 115: "ROUND", 116: "HARD_SWISH", 117: "IF", 118: "WHILE", 119: "NON_MAX_SUPPRESSION_V4", 120: "NON_MAX_SUPPRESSION_V5", 121: "SCATTER_ND", 122: "SELECT_V2", 123: "DENSIFY", 124: "SEGMENT_SUM", 125: "BATCH_MATMUL", 126: "PLACEHOLDER_FOR_GREATER_OP_CODES", 127: "CUMSUM", 128: "CALL_ONCE", 129: "BROADCAST_TO", 130: "RFFT2D", 131: "CONV_3D", 132: "IMAG", 133: "REAL", 134: "COMPLEX_ABS", 135: "HASHTABLE", 136: "HASHTABLE_FIND", 137: "HASHTABLE_IMPORT", 138: "HASHTABLE_SIZE", 139: "REDUCE_ALL", 140: "CONV_3D_TRANSPOSE", 141: "VAR_HANDLE", 142: "READ_VARIABLE", 143: "ASSIGN_VARIABLE", 144: "BROADCAST_ARGS", 145: "RANDOM_STANDARD_NORMAL", 146: "BUCKETIZE", 147: "RANDOM_UNIFORM", 148: "MULTINOMIAL", 149: "GELU", 150: "DYNAMIC_UPDATE_SLICE", 151: "IRFFT2D", 152: "EXP", 153: "PRELU", 154: "MAXIMUM", 155: "ARG_MAX", 156: "ARG_MIN", 157: "GELU", 158: "DYNAMIC_UPDATE_SLICE", 159: "RELU_0_TO_1", 160: "REDUCE_PROD", 161: "RELU_0_TO_1", 162: "REDUCE_PROD"} class FB: def __init__(self, data): self.d = data def u32(self, p): return struct.unpack_from(" vt + vs: return None f = self.u16(off) return None if f == 0 else t + f def deref(self, p): return p + self.u32(p) if p is not None else None def vec(self, p): p = self.deref(p) if p is None: return None return p + 4, self.u32(p) def vtab(self, p): s, n = self.vec(p) return [s + i * 4 + self.u32(s + i * 4) for i in range(n)] def ivec(self, p): s, n = self.vec(p) return [self.i32(s + 4 * i) for i in range(n)] def main(path, n_ops=8): fb = FB(open(path, "rb").read()) root = fb.u32(0) model = root opcodes = fb.vtab(fb.field(model, 1)) subgraphs = fb.vtab(fb.field(model, 2)) sg = subgraphs[0] tensors = fb.vtab(fb.field(sg, 0)) ops = fb.vtab(fb.field(sg, 3)) buffers = fb.vtab(fb.field(model, 4)) def tdesc(i): t = tensors[i] shape = fb.ivec(fb.field(t, 0)) tt = fb.field(t, 1) ttype = TENSOR_TYPE.get(fb.d[tt], fb.d[tt]) if tt else "?" buf = fb.field(t, 2) bi = fb.u32(buf) if buf else 0 name = fb.deref(fb.field(t, 3)) if name: name = fb.d[name + 4:name + 4 + fb.u32(name)].decode("utf-8", "replace") return i, name, shape, ttype, bi def opcode_of(op): idx_f = fb.field(op, 0) # opcode_index 是内联标量 idx = fb.u32(idx_f) if idx_f else 0 oc = opcodes[idx] bc_f = fb.field(oc, 1) bc = fb.i32(bc_f) if bc_f else 0 return BUILTIN.get(bc, bc) for oi, op in enumerate(ops[:n_ops]): ins = fb.ivec(fb.field(op, 1)) outs = fb.ivec(fb.field(op, 2)) print(f"op{oi} {opcode_of(op)}") for i in ins: ti, name, shape, ttype, bi = tdesc(i) extra = "" if ttype in ("INT32", "UINT8", "INT64", "INT16", "UINT16", "INT8") and bi and bi < len(buffers): data = buffers[bi] if data: start = fb.field(data, 0) size = fb.u32(start) if start else 0 if 0 < size <= 64: extra = f" const={list(fb.d[start + 4:start + 4 + size])}" print(f" in {name} {shape} {ttype}{extra}") for i in outs: ti, name, shape, ttype, bi = tdesc(i) print(f" out {name} {shape} {ttype}") if __name__ == "__main__": main(sys.argv[1], int(sys.argv[2]) if len(sys.argv) > 2 else 8)