433 lines
14 KiB
Python
433 lines
14 KiB
Python
|
# 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)
|