wehub-resource-sync
300 lines
11 KiB
Python
300 lines
11 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 that support testing.
|
|
|
|
All functions that can be commonly used by various tests.
|
|
"""
|
|
|
|
import flatbuffers
|
|
from tensorflow.lite.python import schema_py_generated as schema_fb
|
|
|
|
TFLITE_SCHEMA_VERSION = 3
|
|
|
|
|
|
def build_mock_flatbuffer_model():
|
|
"""Creates a flatbuffer containing an example model."""
|
|
builder = flatbuffers.Builder(1024)
|
|
|
|
schema_fb.BufferStart(builder)
|
|
buffer0_offset = schema_fb.BufferEnd(builder)
|
|
|
|
schema_fb.BufferStartDataVector(builder, 12)
|
|
builder.PrependUint8(11)
|
|
builder.PrependUint8(10)
|
|
builder.PrependUint8(9)
|
|
builder.PrependUint8(8)
|
|
builder.PrependUint8(7)
|
|
builder.PrependUint8(6)
|
|
builder.PrependUint8(5)
|
|
builder.PrependUint8(4)
|
|
builder.PrependUint8(3)
|
|
builder.PrependUint8(2)
|
|
builder.PrependUint8(1)
|
|
builder.PrependUint8(0)
|
|
buffer1_data_offset = builder.EndVector()
|
|
schema_fb.BufferStart(builder)
|
|
schema_fb.BufferAddData(builder, buffer1_data_offset)
|
|
buffer1_offset = schema_fb.BufferEnd(builder)
|
|
|
|
schema_fb.BufferStart(builder)
|
|
buffer2_offset = schema_fb.BufferEnd(builder)
|
|
|
|
schema_fb.ModelStartBuffersVector(builder, 3)
|
|
builder.PrependUOffsetTRelative(buffer2_offset)
|
|
builder.PrependUOffsetTRelative(buffer1_offset)
|
|
builder.PrependUOffsetTRelative(buffer0_offset)
|
|
buffers_offset = builder.EndVector()
|
|
|
|
string0_offset = builder.CreateString('input_tensor')
|
|
schema_fb.TensorStartShapeVector(builder, 3)
|
|
builder.PrependInt32(1)
|
|
builder.PrependInt32(2)
|
|
builder.PrependInt32(5)
|
|
shape0_offset = builder.EndVector()
|
|
schema_fb.TensorStart(builder)
|
|
schema_fb.TensorAddName(builder, string0_offset)
|
|
schema_fb.TensorAddShape(builder, shape0_offset)
|
|
schema_fb.TensorAddType(builder, 0)
|
|
schema_fb.TensorAddBuffer(builder, 0)
|
|
tensor0_offset = schema_fb.TensorEnd(builder)
|
|
|
|
schema_fb.QuantizationParametersStartMinVector(builder, 5)
|
|
builder.PrependFloat32(0.5)
|
|
builder.PrependFloat32(2.0)
|
|
builder.PrependFloat32(5.0)
|
|
builder.PrependFloat32(10.0)
|
|
builder.PrependFloat32(20.0)
|
|
quant1_min_offset = builder.EndVector()
|
|
|
|
schema_fb.QuantizationParametersStartMaxVector(builder, 5)
|
|
builder.PrependFloat32(10.0)
|
|
builder.PrependFloat32(20.0)
|
|
builder.PrependFloat32(-50.0)
|
|
builder.PrependFloat32(1.0)
|
|
builder.PrependFloat32(2.0)
|
|
quant1_max_offset = builder.EndVector()
|
|
|
|
schema_fb.QuantizationParametersStartScaleVector(builder, 5)
|
|
builder.PrependFloat32(3.0)
|
|
builder.PrependFloat32(4.0)
|
|
builder.PrependFloat32(5.0)
|
|
builder.PrependFloat32(6.0)
|
|
builder.PrependFloat32(7.0)
|
|
quant1_scale_offset = builder.EndVector()
|
|
|
|
schema_fb.QuantizationParametersStartZeroPointVector(builder, 5)
|
|
builder.PrependInt64(1)
|
|
builder.PrependInt64(2)
|
|
builder.PrependInt64(3)
|
|
builder.PrependInt64(-1)
|
|
builder.PrependInt64(-2)
|
|
quant1_zero_point_offset = builder.EndVector()
|
|
|
|
schema_fb.QuantizationParametersStart(builder)
|
|
schema_fb.QuantizationParametersAddMin(builder, quant1_min_offset)
|
|
schema_fb.QuantizationParametersAddMax(builder, quant1_max_offset)
|
|
schema_fb.QuantizationParametersAddScale(builder, quant1_scale_offset)
|
|
schema_fb.QuantizationParametersAddZeroPoint(builder,
|
|
quant1_zero_point_offset)
|
|
quantization1_offset = schema_fb.QuantizationParametersEnd(builder)
|
|
|
|
string1_offset = builder.CreateString('constant_tensor')
|
|
schema_fb.TensorStartShapeVector(builder, 3)
|
|
builder.PrependInt32(1)
|
|
builder.PrependInt32(2)
|
|
builder.PrependInt32(5)
|
|
shape1_offset = builder.EndVector()
|
|
schema_fb.TensorStart(builder)
|
|
schema_fb.TensorAddName(builder, string1_offset)
|
|
schema_fb.TensorAddShape(builder, shape1_offset)
|
|
schema_fb.TensorAddType(builder, schema_fb.TensorType.UINT8)
|
|
schema_fb.TensorAddBuffer(builder, 1)
|
|
schema_fb.TensorAddQuantization(builder, quantization1_offset)
|
|
tensor1_offset = schema_fb.TensorEnd(builder)
|
|
|
|
string2_offset = builder.CreateString('output_tensor')
|
|
schema_fb.TensorStartShapeVector(builder, 3)
|
|
builder.PrependInt32(1)
|
|
builder.PrependInt32(2)
|
|
builder.PrependInt32(5)
|
|
shape2_offset = builder.EndVector()
|
|
schema_fb.TensorStart(builder)
|
|
schema_fb.TensorAddName(builder, string2_offset)
|
|
schema_fb.TensorAddShape(builder, shape2_offset)
|
|
schema_fb.TensorAddType(builder, 0)
|
|
schema_fb.TensorAddBuffer(builder, 2)
|
|
tensor2_offset = schema_fb.TensorEnd(builder)
|
|
|
|
schema_fb.SubGraphStartTensorsVector(builder, 3)
|
|
builder.PrependUOffsetTRelative(tensor2_offset)
|
|
builder.PrependUOffsetTRelative(tensor1_offset)
|
|
builder.PrependUOffsetTRelative(tensor0_offset)
|
|
tensors_offset = builder.EndVector()
|
|
|
|
schema_fb.SubGraphStartInputsVector(builder, 1)
|
|
builder.PrependInt32(0)
|
|
inputs_offset = builder.EndVector()
|
|
|
|
schema_fb.SubGraphStartOutputsVector(builder, 1)
|
|
builder.PrependInt32(2)
|
|
outputs_offset = builder.EndVector()
|
|
|
|
schema_fb.OperatorCodeStart(builder)
|
|
schema_fb.OperatorCodeAddBuiltinCode(builder, schema_fb.BuiltinOperator.ADD)
|
|
schema_fb.OperatorCodeAddDeprecatedBuiltinCode(builder,
|
|
schema_fb.BuiltinOperator.ADD)
|
|
schema_fb.OperatorCodeAddVersion(builder, 1)
|
|
code0_offset = schema_fb.OperatorCodeEnd(builder)
|
|
|
|
schema_fb.OperatorCodeStart(builder)
|
|
schema_fb.OperatorCodeAddBuiltinCode(builder,
|
|
schema_fb.BuiltinOperator.VAR_HANDLE)
|
|
schema_fb.OperatorCodeAddDeprecatedBuiltinCode(
|
|
builder, schema_fb.BuiltinOperator.PLACEHOLDER_FOR_GREATER_OP_CODES)
|
|
schema_fb.OperatorCodeAddVersion(builder, 1)
|
|
code1_offset = schema_fb.OperatorCodeEnd(builder)
|
|
|
|
schema_fb.ModelStartOperatorCodesVector(builder, 2)
|
|
builder.PrependUOffsetTRelative(code1_offset)
|
|
builder.PrependUOffsetTRelative(code0_offset)
|
|
codes_offset = builder.EndVector()
|
|
|
|
schema_fb.OperatorStartInputsVector(builder, 2)
|
|
builder.PrependInt32(0)
|
|
builder.PrependInt32(1)
|
|
op_inputs_offset = builder.EndVector()
|
|
|
|
schema_fb.OperatorStartOutputsVector(builder, 1)
|
|
builder.PrependInt32(2)
|
|
op_outputs_offset = builder.EndVector()
|
|
|
|
schema_fb.OperatorStart(builder)
|
|
schema_fb.OperatorAddOpcodeIndex(builder, 0)
|
|
schema_fb.OperatorAddInputs(builder, op_inputs_offset)
|
|
schema_fb.OperatorAddOutputs(builder, op_outputs_offset)
|
|
op0_offset = schema_fb.OperatorEnd(builder)
|
|
|
|
shared_name = builder.CreateString('var')
|
|
schema_fb.VarHandleOptionsStart(builder)
|
|
schema_fb.VarHandleOptionsAddSharedName(builder, shared_name)
|
|
var_handle_options_offset = schema_fb.VarHandleOptionsEnd(builder)
|
|
|
|
schema_fb.OperatorStart(builder)
|
|
schema_fb.OperatorAddOpcodeIndex(builder, 1)
|
|
schema_fb.OperatorAddBuiltinOptionsType(
|
|
builder, schema_fb.BuiltinOptions.VarHandleOptions)
|
|
schema_fb.OperatorAddBuiltinOptions(builder, var_handle_options_offset)
|
|
op1_offset = schema_fb.OperatorEnd(builder)
|
|
|
|
schema_fb.OperatorStart(builder)
|
|
schema_fb.OperatorAddBuiltinOptionsType(
|
|
builder, schema_fb.BuiltinOptions.VarHandleOptions)
|
|
schema_fb.OperatorAddBuiltinOptions(builder, var_handle_options_offset)
|
|
op2_offset = schema_fb.OperatorEnd(builder)
|
|
|
|
schema_fb.SubGraphStartOperatorsVector(builder, 3)
|
|
builder.PrependUOffsetTRelative(op2_offset)
|
|
builder.PrependUOffsetTRelative(op1_offset)
|
|
builder.PrependUOffsetTRelative(op0_offset)
|
|
ops_offset = builder.EndVector()
|
|
|
|
string3_offset = builder.CreateString('subgraph_name')
|
|
schema_fb.SubGraphStart(builder)
|
|
schema_fb.SubGraphAddName(builder, string3_offset)
|
|
schema_fb.SubGraphAddTensors(builder, tensors_offset)
|
|
schema_fb.SubGraphAddInputs(builder, inputs_offset)
|
|
schema_fb.SubGraphAddOutputs(builder, outputs_offset)
|
|
schema_fb.SubGraphAddOperators(builder, ops_offset)
|
|
subgraph_offset = schema_fb.SubGraphEnd(builder)
|
|
|
|
schema_fb.ModelStartSubgraphsVector(builder, 1)
|
|
builder.PrependUOffsetTRelative(subgraph_offset)
|
|
subgraphs_offset = builder.EndVector()
|
|
|
|
signature_key = builder.CreateString('my_key')
|
|
input_tensor_string = builder.CreateString('input_tensor')
|
|
output_tensor_string = builder.CreateString('output_tensor')
|
|
|
|
# Signature Inputs
|
|
schema_fb.TensorMapStart(builder)
|
|
schema_fb.TensorMapAddName(builder, input_tensor_string)
|
|
schema_fb.TensorMapAddTensorIndex(builder, 1)
|
|
input_tensor = schema_fb.TensorMapEnd(builder)
|
|
|
|
# Signature Outputs
|
|
schema_fb.TensorMapStart(builder)
|
|
schema_fb.TensorMapAddName(builder, output_tensor_string)
|
|
schema_fb.TensorMapAddTensorIndex(builder, 2)
|
|
output_tensor = schema_fb.TensorMapEnd(builder)
|
|
|
|
schema_fb.SignatureDefStartInputsVector(builder, 1)
|
|
builder.PrependUOffsetTRelative(input_tensor)
|
|
signature_inputs_offset = builder.EndVector()
|
|
schema_fb.SignatureDefStartOutputsVector(builder, 1)
|
|
builder.PrependUOffsetTRelative(output_tensor)
|
|
signature_outputs_offset = builder.EndVector()
|
|
|
|
schema_fb.SignatureDefStart(builder)
|
|
schema_fb.SignatureDefAddSignatureKey(builder, signature_key)
|
|
schema_fb.SignatureDefAddInputs(builder, signature_inputs_offset)
|
|
schema_fb.SignatureDefAddOutputs(builder, signature_outputs_offset)
|
|
signature_offset = schema_fb.SignatureDefEnd(builder)
|
|
schema_fb.ModelStartSignatureDefsVector(builder, 1)
|
|
builder.PrependUOffsetTRelative(signature_offset)
|
|
signature_defs_offset = builder.EndVector()
|
|
|
|
string4_offset = builder.CreateString('model_description')
|
|
schema_fb.ModelStart(builder)
|
|
schema_fb.ModelAddVersion(builder, TFLITE_SCHEMA_VERSION)
|
|
schema_fb.ModelAddOperatorCodes(builder, codes_offset)
|
|
schema_fb.ModelAddSubgraphs(builder, subgraphs_offset)
|
|
schema_fb.ModelAddDescription(builder, string4_offset)
|
|
schema_fb.ModelAddBuffers(builder, buffers_offset)
|
|
schema_fb.ModelAddSignatureDefs(builder, signature_defs_offset)
|
|
model_offset = schema_fb.ModelEnd(builder)
|
|
builder.Finish(model_offset)
|
|
model = builder.Output()
|
|
|
|
return model
|
|
|
|
|
|
def build_operator_with_options() -> schema_fb.Operator:
|
|
"""Builds an operator with the given options."""
|
|
builder = flatbuffers.Builder(1024)
|
|
schema_fb.StableHLOCompositeOptionsStart(builder)
|
|
schema_fb.StableHLOCompositeOptionsAddDecompositionSubgraphIndex(builder, 10)
|
|
opts = schema_fb.StableHLOCompositeOptionsEnd(builder)
|
|
schema_fb.OperatorStart(builder)
|
|
schema_fb.OperatorAddBuiltinOptions2(builder, opts)
|
|
schema_fb.OperatorAddBuiltinOptions2Type(
|
|
builder, schema_fb.BuiltinOptions2.StableHLOCompositeOptions
|
|
)
|
|
op_offset = schema_fb.OperatorEnd(builder)
|
|
builder.Finish(op_offset)
|
|
return schema_fb.Operator.GetRootAs(builder.Output())
|
|
|
|
|
|
def load_model_from_flatbuffer(flatbuffer_model):
|
|
"""Loads a model as a python object from a flatbuffer model."""
|
|
model = schema_fb.Model.GetRootAsModel(flatbuffer_model, 0)
|
|
model = schema_fb.ModelT.InitFromObj(model)
|
|
return model
|
|
|
|
|
|
def build_mock_model():
|
|
"""Creates an object containing an example model."""
|
|
model = build_mock_flatbuffer_model()
|
|
return load_model_from_flatbuffer(model)
|