tensorflow/lite/tools/flatbuffer_utils.py
# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility functions for FlatBuffers.
All functions that are commonly used to work with FlatBuffers.
Refer to the tensorflow lite flatbuffer schema here:
tensorflow/lite/schema/schema.fbs
"""
import copy
import random
import re
import struct
import sys
import flatbuffers
from tensorflow.lite.python import schema_py_generated as schema_fb
from tensorflow.lite.python import schema_util
from tensorflow.python.platform import gfile
_TFLITE_FILE_IDENTIFIER = b'TFL3'
def convert_bytearray_to_object(model_bytearray):
"""Converts a tflite model from a bytearray to an object for parsing."""
model_object = schema_fb.Model.GetRootAsModel(model_bytearray, 0)
return schema_fb.ModelT.InitFromObj(model_object)
def read_model(input_tflite_file):
"""Reads a tflite model as a python object.
Args:
input_tflite_file: Full path name to the input tflite file
Raises:
RuntimeError: If input_tflite_file path is invalid.
IOError: If input_tflite_file cannot be opened.
Returns:
A python object corresponding to the input tflite file.
"""
if not gfile.Exists(input_tflite_file):
raise RuntimeError('Input file not found at %r\n' % input_tflite_file)
with gfile.GFile(input_tflite_file, 'rb') as input_file_handle:
model_bytearray = bytearray(input_file_handle.read())
model = convert_bytearray_to_object(model_bytearray)
if sys.byteorder == 'big':
byte_swap_tflite_model_obj(model, 'little', 'big')
return model
def read_model_with_mutable_tensors(input_tflite_file):
"""Reads a tflite model as a python object with mutable tensors.
Similar to read_model() with the addition that the returned object has
mutable tensors (read_model() returns an object with immutable tensors).
NOTE: This API only works for TFLite generated with
_experimental_use_buffer_offset=false
Args:
input_tflite_file: Full path name to the input tflite file
Raises:
RuntimeError: If input_tflite_file path is invalid.
IOError: If input_tflite_file cannot be opened.
Returns:
A mutable python object corresponding to the input tflite file.
"""
return copy.deepcopy(read_model(input_tflite_file))
def convert_object_to_bytearray(model_object, extra_buffer=b''):
"""Converts a tflite model from an object to a immutable bytearray."""
# Initial size of the buffer, which will grow automatically if needed
builder = flatbuffers.Builder(1024)
model_offset = model_object.Pack(builder)
builder.Finish(model_offset, file_identifier=_TFLITE_FILE_IDENTIFIER)
model_bytearray = bytes(builder.Output())
model_bytearray = model_bytearray + extra_buffer
return model_bytearray
def write_model(model_object, output_tflite_file):
"""Writes the tflite model, a python object, into the output file.
NOTE: This API only works for TFLite generated with
_experimental_use_buffer_offset=false
Args:
model_object: A tflite model as a python object
output_tflite_file: Full path name to the output tflite file.
Raises:
IOError: If output_tflite_file path is invalid or cannot be opened.
"""
if sys.byteorder == 'big':
model_object = copy.deepcopy(model_object)
byte_swap_tflite_model_obj(model_object, 'big', 'little')
model_bytearray = convert_object_to_bytearray(model_object)
with gfile.GFile(output_tflite_file, 'wb') as output_file_handle:
output_file_handle.write(model_bytearray)
def strip_strings(model):
"""Strips all nonessential strings from the model to reduce model size.
We remove the following strings:
(find strings by searching ":string" in the tensorflow lite flatbuffer schema)
1. Model description
2. SubGraph name
3. Tensor names
We retain OperatorCode custom_code and Metadata name.
Args:
model: The model from which to remove nonessential strings.
"""
model.description = None
for subgraph in model.subgraphs:
subgraph.name = None
for tensor in subgraph.tensors:
tensor.name = None
# We clear all signature_def structure, since without names it is useless.
model.signatureDefs = None
def type_to_name(tensor_type):
"""Converts a numerical enum to a readable tensor type."""
for name, value in schema_fb.TensorType.__dict__.items():
if value == tensor_type:
return name
return None
def randomize_weights(model, random_seed=0, buffers_to_skip=None):
"""Randomize weights in a model.
Args:
model: The model in which to randomize weights.
random_seed: The input to the random number generator (default value is 0).
buffers_to_skip: The list of buffer indices to skip. The weights in these
buffers are left unmodified.
"""
# The input to the random seed generator. The default value is 0.
random.seed(random_seed)
# Parse model buffers which store the model weights
buffers = model.buffers
buffer_ids = range(1, len(buffers)) # ignore index 0 as it's always None
if buffers_to_skip is not None:
buffer_ids = [idx for idx in buffer_ids if idx not in buffers_to_skip]
buffer_types = {}
for graph in model.subgraphs:
for op in graph.operators:
if op.inputs is None:
break
for input_idx in op.inputs:
tensor = graph.tensors[input_idx]
buffer_types[tensor.buffer] = type_to_name(tensor.type)
for i in buffer_ids:
buffer_i_data = buffers[i].data
buffer_i_size = 0 if buffer_i_data is None else buffer_i_data.size
if buffer_i_size == 0:
continue
# Raw data buffers are of type ubyte (or uint8) whose values lie in the
# range [0, 255]. Those ubytes (or unint8s) are the underlying
# representation of each datatype. For example, a bias tensor of type
# int32 appears as a buffer 4 times it's length of type ubyte (or uint8).
# For floats, we need to generate a valid float and then pack it into
# the raw bytes in place.
buffer_type = buffer_types.get(i, 'INT8')
if buffer_type.startswith('FLOAT'):
format_code = 'e' if buffer_type == 'FLOAT16' else 'f'
for offset in range(0, buffer_i_size, struct.calcsize(format_code)):
value = random.uniform(-0.5, 0.5) # See http://b/152324470#comment2
struct.pack_into(format_code, buffer_i_data, offset, value)
else:
for j in range(buffer_i_size):
buffer_i_data[j] = random.randint(0, 255)
def rename_custom_ops(model, map_custom_op_renames):
"""Rename custom ops so they use the same naming style as builtin ops.
Args:
model: The input tflite model.
map_custom_op_renames: A mapping from old to new custom op names.
"""
for op_code in model.operatorCodes:
if op_code.customCode:
op_code_str = op_code.customCode.decode('ascii')
if op_code_str in map_custom_op_renames:
op_code.customCode = map_custom_op_renames[op_code_str].encode('ascii')
def opcode_to_name(model, op_code):
"""Converts a TFLite op_code to the human readable name.
Args:
model: The input tflite model.
op_code: The op_code to resolve to a readable name.
Returns:
A string containing the human readable op name, or None if not resolvable.
"""
op = model.operatorCodes[op_code]
code = max(op.builtinCode, op.deprecatedBuiltinCode)
for name, value in vars(schema_fb.BuiltinOperator).items():
if value == code:
return name
return None
def xxd_output_to_bytes(input_cc_file):
"""Converts xxd output C++ source file to bytes (immutable).
Args:
input_cc_file: Full path name to th C++ source file dumped by xxd
Raises:
RuntimeError: If input_cc_file path is invalid.
IOError: If input_cc_file cannot be opened.
Returns:
A bytearray corresponding to the input cc file array.
"""
# Match hex values in the string with comma as separator
pattern = re.compile(r'\W*(0x[0-9a-fA-F,x ]+).*')
model_bytearray = bytearray()
with open(input_cc_file) as file_handle:
for line in file_handle:
values_match = pattern.match(line)
if values_match is None:
continue
# Match in the parentheses (hex array only)
list_text = values_match.group(1)
# Extract hex values (text) from the line
# e.g. 0x1c, 0x00, 0x00, 0x00, 0x54, 0x46, 0x4c,
values_text = filter(None, list_text.split(','))
# Convert to hex
values = [int(x, base=16) for x in values_text]
model_bytearray.extend(values)
return bytes(model_bytearray)
def xxd_output_to_object(input_cc_file):
"""Converts xxd output C++ source file to object.
Args:
input_cc_file: Full path name to th C++ source file dumped by xxd
Raises:
RuntimeError: If input_cc_file path is invalid.
IOError: If input_cc_file cannot be opened.
Returns:
A python object corresponding to the input tflite file.
"""
model_bytes = xxd_output_to_bytes(input_cc_file)
return convert_bytearray_to_object(model_bytes)
def byte_swap_buffer_content(buffer, chunksize, from_endiness, to_endiness):
"""Helper function for byte-swapping the buffers field."""
to_swap = [
buffer.data[i : i + chunksize]
for i in range(0, len(buffer.data), chunksize)
]
buffer.data = b''.join(
[
int.from_bytes(byteswap, from_endiness).to_bytes(
chunksize, to_endiness
)
for byteswap in to_swap
]
)
def byte_swap_string_content(buffer, from_endiness, to_endiness):
"""Helper function for byte-swapping the string buffer.
Args:
buffer: TFLite string buffer of from_endiness format.
from_endiness: The original endianness format of the string buffer.
to_endiness: The destined endianness format of the string buffer.
"""
num_of_strings = int.from_bytes(buffer.data[0:4], from_endiness)
string_content = bytearray(buffer.data[4 * (num_of_strings + 2) :])
prefix_data = b''.join(
[
int.from_bytes(buffer.data[i : i + 4], from_endiness).to_bytes(
4, to_endiness
)
for i in range(0, (num_of_strings + 1) * 4 + 1, 4)
]
)
buffer.data = prefix_data + string_content
def byte_swap_tflite_model_obj(model, from_endiness, to_endiness):
"""Byte swaps the buffers field in a TFLite model.
Args:
model: TFLite model object of from_endiness format.
from_endiness: The original endianness format of the buffers in model.
to_endiness: The destined endianness format of the buffers in model.
"""
if model is None:
return
# Get all the constant buffers, byte swapping them as per their data types
buffer_swapped = []
types_of_16_bits = [
schema_fb.TensorType.FLOAT16,
schema_fb.TensorType.INT16,
schema_fb.TensorType.UINT16,
]
types_of_32_bits = [
schema_fb.TensorType.FLOAT32,
schema_fb.TensorType.INT32,
schema_fb.TensorType.COMPLEX64,
schema_fb.TensorType.UINT32,
]
types_of_64_bits = [
schema_fb.TensorType.INT64,
schema_fb.TensorType.FLOAT64,
schema_fb.TensorType.COMPLEX128,
schema_fb.TensorType.UINT64,
]
for subgraph in model.subgraphs:
for tensor in subgraph.tensors:
if (
tensor.buffer > 0
and tensor.buffer < len(model.buffers)
and tensor.buffer not in buffer_swapped
and model.buffers[tensor.buffer].data is not None
):
if tensor.type == schema_fb.TensorType.STRING:
byte_swap_string_content(
model.buffers[tensor.buffer], from_endiness, to_endiness
)
elif tensor.type in types_of_16_bits:
byte_swap_buffer_content(
model.buffers[tensor.buffer], 2, from_endiness, to_endiness
)
elif tensor.type in types_of_32_bits:
byte_swap_buffer_content(
model.buffers[tensor.buffer], 4, from_endiness, to_endiness
)
elif tensor.type in types_of_64_bits:
byte_swap_buffer_content(
model.buffers[tensor.buffer], 8, from_endiness, to_endiness
)
else:
continue
buffer_swapped.append(tensor.buffer)
def byte_swap_tflite_buffer(tflite_model, from_endiness, to_endiness):
"""Generates a new model byte array after byte swapping its buffers field.
Args:
tflite_model: TFLite flatbuffer in a byte array.
from_endiness: The original endianness format of the buffers in
tflite_model.
to_endiness: The destined endianness format of the buffers in tflite_model.
Returns:
TFLite flatbuffer in a byte array, after being byte swapped to to_endiness
format.
"""
if tflite_model is None:
return None
# Load TFLite Flatbuffer byte array into an object.
model = convert_bytearray_to_object(tflite_model)
# Byte swapping the constant buffers as per their data types
byte_swap_tflite_model_obj(model, from_endiness, to_endiness)
# Return a TFLite flatbuffer as a byte array.
return convert_object_to_bytearray(model)
def count_resource_variables(model):
"""Calculates the number of unique resource variables in a model.
Args:
model: the input tflite model, either as bytearray or object.
Returns:
An integer number representing the number of unique resource variables.
"""
if not isinstance(model, schema_fb.ModelT):
model = convert_bytearray_to_object(model)
unique_shared_names = set()
for subgraph in model.subgraphs:
if subgraph.operators is None:
continue
for op in subgraph.operators:
builtin_code = schema_util.get_builtin_code_from_operator_code(
model.operatorCodes[op.opcodeIndex]
)
if builtin_code == schema_fb.BuiltinOperator.VAR_HANDLE:
unique_shared_names.add(op.builtinOptions.sharedName)
return len(unique_shared_names)