wehub-resource-sync
5586 lines
202 KiB
Python
5586 lines
202 KiB
Python
# Copyright 2019 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.
|
|
# ==============================================================================
|
|
"""Tests for lite.py functionality related to TensorFlow 2.0."""
|
|
|
|
import ctypes
|
|
import functools
|
|
import itertools
|
|
import os
|
|
import sys
|
|
|
|
from absl.testing import parameterized
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
# Force loaded shared object symbols to be globally visible. This is needed so
|
|
# that the interpreter_wrapper, in one .so file, can see the test_registerer,
|
|
# in a different .so file. Note that this may already be set by default.
|
|
# pylint: disable=g-import-not-at-top
|
|
if hasattr(sys, 'setdlopenflags') and hasattr(sys, 'getdlopenflags'):
|
|
sys.setdlopenflags(sys.getdlopenflags() | ctypes.RTLD_GLOBAL)
|
|
|
|
from tensorflow.compiler.mlir.quantization.stablehlo import quantization_options_pb2 as quant_opts_pb2
|
|
from tensorflow.lite.python import conversion_metadata_schema_py_generated as metadata_fb
|
|
from tensorflow.lite.python import convert
|
|
from tensorflow.lite.python import interpreter
|
|
from tensorflow.lite.python import lite
|
|
from tensorflow.lite.python import lite_v2_test_util
|
|
from tensorflow.lite.python import schema_py_generated as schema_fb
|
|
from tensorflow.lite.python import test_util as tflite_test_util
|
|
from tensorflow.lite.python import util
|
|
from tensorflow.lite.python.testdata import _pywrap_test_registerer as test_registerer
|
|
from tensorflow.lite.python.testdata import double_op
|
|
from tensorflow.lite.tools import flatbuffer_utils
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import test_util
|
|
from tensorflow.python.framework import versions
|
|
from tensorflow.python.lib.io import file_io
|
|
from tensorflow.python.ops import map_ops
|
|
from tensorflow.python.ops import rnn
|
|
from tensorflow.python.platform import resource_loader
|
|
from tensorflow.python.platform import test
|
|
from tensorflow.python.saved_model import loader_impl
|
|
from tensorflow.python.saved_model import save
|
|
from tensorflow.python.saved_model import save_options
|
|
from tensorflow.python.saved_model import saved_model
|
|
from tensorflow.python.trackable import autotrackable
|
|
|
|
# Type alias for preset quantization method protobuf enums.
|
|
_PresetQuantizationMethod = quant_opts_pb2.PresetQuantizationMethod.PresetMethod
|
|
|
|
# Only run jax related tests when we can import jax.
|
|
DISABLE_JAX_TEST = False
|
|
try:
|
|
import jax
|
|
from jax import numpy as jnp
|
|
except ImportError:
|
|
DISABLE_JAX_TEST = True
|
|
# pylint: enable=g-import-not-at-top
|
|
|
|
|
|
class FromConcreteFunctionTest(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testTypeInvalid(self):
|
|
root = self._getSimpleVariableModel()
|
|
with self.assertRaises(ValueError) as error:
|
|
_ = lite.TFLiteConverterV2.from_concrete_functions([root.f], root)
|
|
self.assertIn('call get_concrete_function', str(error.exception))
|
|
|
|
@test_util.run_v2_only
|
|
def testFloat(self):
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check output value from converted model.
|
|
expected_value = root.f(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
@parameterized.named_parameters(
|
|
('_INT8InputOutput', dtypes.int8),
|
|
('_UINT8InputOutput', dtypes.uint8),
|
|
('_INT16InputOutput', dtypes.int16),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testInvalidFloat(self, inference_input_output_type):
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
with self.assertRaises(ValueError) as error:
|
|
converter.inference_input_type = inference_input_output_type
|
|
converter.inference_output_type = inference_input_output_type
|
|
converter.convert()
|
|
self.assertEqual(
|
|
'The inference_input_type and inference_output_type '
|
|
'must be tf.float32.',
|
|
str(error.exception),
|
|
)
|
|
|
|
@test_util.run_v2_only
|
|
def testScalarInput(self):
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.f(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testStringInput(self):
|
|
class Model(tf.Module):
|
|
|
|
@tf.function
|
|
def __call__(self, x):
|
|
return x
|
|
|
|
root = Model()
|
|
concrete_func = root.__call__.get_concrete_function(
|
|
tf.constant([str(x) for x in range(11)])
|
|
)
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
tflite_model = converter.convert()
|
|
input_data = tf.constant(
|
|
[str(x) for x in range(11)], shape=(11,), dtype=tf.dtypes.string
|
|
)
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
interp.allocate_tensors()
|
|
my_signature = interp.get_signature_runner()
|
|
|
|
with self.assertRaises(ValueError) as error:
|
|
_ = my_signature(x=input_data)
|
|
self.assertIn(
|
|
'Passed in value type is not a numpy array, got type ',
|
|
str(error.exception),
|
|
)
|
|
|
|
@test_util.run_v2_only
|
|
def testModelWithoutInputs(self):
|
|
def _get_random_number_gen():
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@tf.function(input_signature=[])
|
|
def func():
|
|
return tf.random.uniform(shape=[1], dtype=tf.float32)
|
|
|
|
root.f = func
|
|
to_save = root.f.get_concrete_function()
|
|
return (root, to_save)
|
|
|
|
# Model with no input
|
|
root, concrete_func = _get_random_number_gen()
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
@test_util.run_v2_only
|
|
def testMultiFunctionModel(self):
|
|
"""Convert a single model in a multi-functional model."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.add.get_concrete_function(input_data)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.add(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testConvertMultipleFunctions(self):
|
|
"""Convert multiple functions in a multi-functional model."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
add_func = root.add.get_concrete_function(input_data)
|
|
sub_func = root.sub.get_concrete_function(input_data)
|
|
|
|
# Try converting multiple functions.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[add_func, sub_func], root
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check signatures are valid from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 2)
|
|
self.assertEqual(list(signature_defs.keys()), ['add', 'sub'])
|
|
self.assertLen(signature_defs.values(), 2)
|
|
self.assertEqual(list(signature_defs['add'].keys()), ['inputs', 'outputs'])
|
|
self.assertCountEqual(signature_defs['add']['inputs'], ['x'])
|
|
self.assertEqual(list(signature_defs['add']['outputs']), ['output_0'])
|
|
self.assertEqual(list(signature_defs['sub'].keys()), ['inputs', 'outputs'])
|
|
self.assertCountEqual(signature_defs['sub']['inputs'], ['x'])
|
|
self.assertEqual(list(signature_defs['sub']['outputs']), ['output_0'])
|
|
|
|
# Verify the Signature runner executions.
|
|
add_signature_runner = interp.get_signature_runner('add')
|
|
add_output = add_signature_runner(x=input_data)
|
|
self.assertEqual(add_output['output_0'], 3)
|
|
input_details = add_signature_runner.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual('add_x:0', input_details['x']['name'])
|
|
self.assertEqual(np.float32, input_details['x']['dtype'])
|
|
self.assertTrue(([1] == input_details['x']['shape']).all())
|
|
self.assertEqual((0.0, 0), input_details['x']['quantization'])
|
|
|
|
sub_signature_runner = interp.get_signature_runner('sub')
|
|
sub_output = sub_signature_runner(x=input_data)
|
|
self.assertEqual(sub_output['output_0'], -2)
|
|
output_details = sub_signature_runner.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertEqual(
|
|
'StatefulPartitionedCall_1:0', output_details['output_0']['name']
|
|
)
|
|
self.assertEqual(np.float32, output_details['output_0']['dtype'])
|
|
self.assertTrue(([1] == output_details['output_0']['shape']).all())
|
|
self.assertEqual((0.0, 0), output_details['output_0']['quantization'])
|
|
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(metadata.environment.apiVersion, 2)
|
|
self.assertEqual(
|
|
metadata.environment.modelType,
|
|
metadata_fb.ModelType.TF_CONCRETE_FUNCTIONS,
|
|
)
|
|
self.assertAllEqual([], metadata.options.modelOptimizationModes)
|
|
|
|
def _getIntegerQuantizeModel(self, num_filters=16):
|
|
np.random.seed(0)
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[1, 5, 5, 3], dtype=tf.float32)]
|
|
)
|
|
def func(inp):
|
|
conv = tf.nn.conv2d(
|
|
inp,
|
|
tf.ones([3, 3, 3, num_filters]),
|
|
strides=[1, 1, 1, 1],
|
|
padding='SAME',
|
|
)
|
|
output = tf.nn.relu(conv, name='output')
|
|
return output
|
|
|
|
def calibration_gen():
|
|
for _ in range(5):
|
|
yield [np.random.uniform(-1, 1, size=(1, 5, 5, 3)).astype(np.float32)]
|
|
|
|
root.f = func
|
|
to_save = root.f.get_concrete_function()
|
|
return (root, to_save, calibration_gen)
|
|
|
|
@parameterized.named_parameters(
|
|
('EnableMlirQuantizer', True), # enable mlir quantizer
|
|
('DisableMlirQuantizer', False),
|
|
) # disable mlir quantizer
|
|
def testPostTrainingCalibrateAndQuantize(self, mlir_quantizer):
|
|
root, func, calibration_gen = self._getIntegerQuantizeModel()
|
|
|
|
# Convert float model.
|
|
float_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
float_tflite_model = float_converter.convert()
|
|
self.assertIsNotNone(float_tflite_model)
|
|
|
|
# Convert quantized model.
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
quantized_converter.experimental_new_quantizer = mlir_quantizer
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(quantized_tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(
|
|
metadata.environment.tensorflowVersion.decode('utf-8'),
|
|
versions.__version__,
|
|
)
|
|
self.assertEqual(metadata.environment.apiVersion, 2)
|
|
self.assertEqual(
|
|
metadata.environment.modelType,
|
|
metadata_fb.ModelType.TF_CONCRETE_FUNCTIONS,
|
|
)
|
|
self.assertEqual(metadata.options.allowCustomOps, False)
|
|
self.assertEqual(metadata.options.enableSelectTfOps, False)
|
|
self.assertEqual(metadata.options.forceSelectTfOps, False)
|
|
self.assertAllEqual(
|
|
[metadata_fb.ModelOptimizationMode.PTQ_FULL_INTEGER],
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
|
|
# The default input and output types should be float.
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual(np.float32, input_details[0]['dtype'])
|
|
output_details = interp.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertEqual(np.float32, output_details[0]['dtype'])
|
|
|
|
# Ensure that the quantized weights tflite model is smaller.
|
|
self.assertLess(len(quantized_tflite_model), len(float_tflite_model))
|
|
|
|
@parameterized.named_parameters(
|
|
('_INT8InputOutput', dtypes.int8),
|
|
('_UINT8InputOutput', dtypes.uint8),
|
|
('_INT16InputOutput', dtypes.int16),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testInvalidPostTrainingDynamicRangeQuantization(
|
|
self, inference_input_output_type
|
|
):
|
|
root, func, _ = self._getIntegerQuantizeModel()
|
|
|
|
# Convert float model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions([func], root)
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
# Convert quantized model.
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
with self.assertRaises(ValueError) as error:
|
|
quantized_converter.inference_input_type = inference_input_output_type
|
|
quantized_converter.inference_output_type = inference_input_output_type
|
|
quantized_converter.convert()
|
|
self.assertEqual(
|
|
'The inference_input_type and inference_output_type '
|
|
'must be tf.float32.',
|
|
str(error.exception),
|
|
)
|
|
|
|
def _createV2QATSavedModelWithFloatOpsAtEnd(self):
|
|
"""Create a simple QAT SavedModel that includes float ops at the end."""
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'qat_float_ops_at_end')
|
|
input_tensor = tf.keras.layers.Input((32, 32, 128))
|
|
|
|
class _FakeQuantArgsLayer(tf.keras.layers.Layer):
|
|
"""A fake quantization layer with fake_quant_with_min_max_args.
|
|
|
|
Keras 3 requires wrapping the tf function inside Keras layer.
|
|
"""
|
|
|
|
def call(self, x):
|
|
return tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
|
|
|
|
x = _FakeQuantArgsLayer()(input_tensor)
|
|
x = tf.keras.layers.Conv2D(1, (3, 3), bias_initializer='ones')(x)
|
|
x = _FakeQuantArgsLayer()(x)
|
|
# Exclude the quantization of the following Dense layer by not putting
|
|
# fake quant layer after the dense layer.
|
|
output_tensor = tf.keras.layers.Dense(
|
|
1, activation='sigmoid', bias_initializer='ones'
|
|
)(x)
|
|
model = tf.keras.Model(input_tensor, output_tensor)
|
|
model.save(saved_model_dir)
|
|
return saved_model_dir
|
|
|
|
def testQuantizationRemovesQDQsForFloatIOInQAT(self):
|
|
saved_model_dir = self._createV2QATSavedModelWithFloatOpsAtEnd()
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_model = converter.convert()
|
|
|
|
# Because assertions on the model later, we opt out applying default TFLite
|
|
# delegates (i.e. the XNNPACK delegate).
|
|
interp = interpreter.Interpreter(
|
|
model_content=quantized_model,
|
|
experimental_op_resolver_type=interpreter.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES,
|
|
)
|
|
interp.allocate_tensors()
|
|
# The model should have LOGISTIC op, instead of DEQUANTIZE op.
|
|
op_details = interp._get_ops_details()
|
|
self.assertEqual(op_details[len(op_details) - 1]['op_name'], 'LOGISTIC')
|
|
|
|
@parameterized.named_parameters(
|
|
('EnableMlirQuantizer', True), # enable mlir quantizer
|
|
('DisableMlirQuantizer', False),
|
|
) # disable mlir quantizer
|
|
def testQuantizationRemovesQDQsForFloatIO(self, mlir_quantizer):
|
|
func, calibration_gen = self._getCeilModel()
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func.get_concrete_function()]
|
|
)
|
|
converter.representative_dataset = calibration_gen
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.experimental_new_quantizer = mlir_quantizer
|
|
quantized_model = converter.convert()
|
|
|
|
# Because assertions on the model later, we opt out applying default TFLite
|
|
# delegates (i.e. the XNNPACK delegate).
|
|
interp = interpreter.Interpreter(
|
|
model_content=quantized_model,
|
|
experimental_op_resolver_type=interpreter.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES,
|
|
)
|
|
interp.allocate_tensors()
|
|
# The model should have only one sqrt op.
|
|
op_details = interp._get_ops_details()
|
|
self.assertLen(op_details, 1)
|
|
self.assertEqual(op_details[0]['op_name'], 'CEIL')
|
|
|
|
@parameterized.named_parameters(
|
|
('_Default', False, False, dtypes.float32),
|
|
('_INT8InputOutput', False, False, dtypes.int8),
|
|
('_UINT8InputOutput', False, False, dtypes.uint8),
|
|
('_INT16Quantize', False, True, dtypes.float32),
|
|
('_INT16Quantize_INT16InputOutput', False, True, dtypes.int16),
|
|
('_IntOnly', True, False, dtypes.float32),
|
|
('_IntOnly_INT8InputOutput', True, False, dtypes.int8),
|
|
('_IntOnly_UINT8InputOutput', True, False, dtypes.uint8),
|
|
('_IntOnly_INT16Quantize', True, True, dtypes.float32),
|
|
('_IntOnly_INT16Quantize_INT16InputOutput', True, True, dtypes.int16),
|
|
)
|
|
def testIntegerQuantization(
|
|
self, is_int_only, is_int16_quantize, inference_input_output_type
|
|
):
|
|
root, func, calibration_gen = self._getIntegerQuantizeModel()
|
|
|
|
# Convert float model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions([func], root)
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
# Convert quantized model.
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
if is_int_only:
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
|
|
]
|
|
else:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8
|
|
]
|
|
else:
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
]
|
|
quantized_converter.inference_input_type = inference_input_output_type
|
|
quantized_converter.inference_output_type = inference_input_output_type
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(quantized_tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
expected_opt_options = [metadata_fb.ModelOptimizationMode.PTQ_FULL_INTEGER]
|
|
if is_int16_quantize:
|
|
expected_opt_options = [metadata_fb.ModelOptimizationMode.PTQ_INT16]
|
|
self.assertAllEqual(
|
|
expected_opt_options, metadata.options.modelOptimizationModes
|
|
)
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual(
|
|
inference_input_output_type.as_numpy_dtype, input_details[0]['dtype']
|
|
)
|
|
output_details = interp.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertEqual(
|
|
inference_input_output_type.as_numpy_dtype, output_details[0]['dtype']
|
|
)
|
|
|
|
# Ensure that the quantized tflite model is smaller.
|
|
self.assertLess(len(quantized_tflite_model), len(tflite_model))
|
|
|
|
@parameterized.named_parameters(('_INT16Quantize_INT8InputOutput', True))
|
|
def testInvalidIntegerQuantization(self, is_int16_quantize):
|
|
root, func, calibration_gen = self._getIntegerQuantizeModel()
|
|
|
|
# Convert quantized model.
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
]
|
|
with self.assertRaises(ValueError) as error:
|
|
quantized_converter.inference_input_type = dtypes.int8
|
|
quantized_converter.inference_output_type = dtypes.int8
|
|
quantized_converter.convert()
|
|
self.assertEqual(
|
|
'The inference_input_type and inference_output_type '
|
|
"must be in ['tf.float32', 'tf.int16'].",
|
|
str(error.exception),
|
|
)
|
|
|
|
def testCalibrateAndQuantizeBuiltinInt16(self):
|
|
root, func, calibration_gen = self._getIntegerQuantizeModel()
|
|
|
|
# Convert float model.
|
|
float_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
float_tflite_model = float_converter.convert()
|
|
self.assertIsNotNone(float_tflite_model)
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions([func], root)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
|
|
]
|
|
converter.representative_dataset = calibration_gen
|
|
quantized_tflite_model = converter.convert()
|
|
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
# The default input and output types should be float.
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual(np.float32, input_details[0]['dtype'])
|
|
output_details = interp.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertEqual(np.float32, output_details[0]['dtype'])
|
|
|
|
# The weights tensor should be quantized to 8 bits,
|
|
# the bias tensor should be 32 bits to utilize optimized kernels,
|
|
# and the activations should be 16 bits.
|
|
tensor_details = interp.get_tensor_details()
|
|
self.assertEqual(np.int8, tensor_details[2]['dtype'])
|
|
self.assertEqual(np.int64, tensor_details[1]['dtype'])
|
|
self.assertEqual(np.int16, tensor_details[0]['dtype'])
|
|
self.assertEqual(np.int16, tensor_details[3]['dtype'])
|
|
|
|
# Ensure that the quantized weights tflite model is smaller.
|
|
self.assertLess(len(quantized_tflite_model), len(float_tflite_model))
|
|
|
|
@test_util.run_v2_only
|
|
def testSignatureDefs(self):
|
|
"""Test converting SignatureDef is correct and uses SignatureDef API."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
add_func = root.add.get_concrete_function(input_data)
|
|
|
|
converter = lite.TFLiteConverterV2([add_func], trackable_obj=root)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = add_func(input_data)
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
results = self._evaluateTFLiteModelUsingSignatureDef(
|
|
tflite_model, 'serving_default', {'x': input_data}
|
|
)
|
|
self.assertLen(list(results.keys()), 1)
|
|
self.assertStartsWith(list(results.keys())[0], 'output')
|
|
self.assertAllClose(
|
|
expected_value.numpy(),
|
|
results[signature_defs['serving_default']['outputs'][0]],
|
|
)
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 1)
|
|
self.assertEqual(list(signature_defs.keys()), ['serving_default'])
|
|
self.assertLen(signature_defs.values(), 1)
|
|
self.assertEqual(
|
|
list(signature_defs['serving_default'].keys()), ['inputs', 'outputs']
|
|
)
|
|
self.assertCountEqual(signature_defs['serving_default']['inputs'], ['x'])
|
|
self.assertLen(list(signature_defs['serving_default']['outputs']), 1)
|
|
self.assertStartsWith(
|
|
list(signature_defs['serving_default']['outputs'])[0], 'output'
|
|
)
|
|
|
|
@test_util.run_v2_only
|
|
def testNoSignatureDefsWhenTrackingObjIsNone(self):
|
|
"""Test converting SignatureDef is correct and uses SignatureDef API."""
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], None
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
# Verify that there is no SignatureDef structure found.
|
|
self.assertEmpty(signature_defs)
|
|
|
|
@test_util.run_v2_only
|
|
def testNoSignatureDefsWhenInvalidTrackingObjIsGiven(self):
|
|
"""Test converting SignatureDef is correct and uses SignatureDef API."""
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], trackable_obj=autotrackable.AutoTrackable()
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
# Verify that there is no SignatureDef structure found.
|
|
self.assertEmpty(signature_defs)
|
|
|
|
@test_util.run_v2_only
|
|
def testTrackbleObject(self):
|
|
"""Test converting with trackable objects."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
add_func = root.add.get_concrete_function(input_data)
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[add_func], trackable_obj=root
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = add_func(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
def _getTrainingTimeQuantizedModel(self):
|
|
class QLinear(tf.keras.layers.Layer):
|
|
|
|
def __init__(self, units=3, **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.units = units
|
|
|
|
def build(self, input_shape):
|
|
self.w = self.add_weight(
|
|
'weight',
|
|
shape=(input_shape[-1], self.units),
|
|
initializer='random_normal',
|
|
trainable=True,
|
|
)
|
|
self.min_var = self.add_weight(
|
|
'min',
|
|
initializer=tf.keras.initializers.Constant(-6.0),
|
|
trainable=False,
|
|
)
|
|
self.max_var = self.add_weight(
|
|
'max',
|
|
initializer=tf.keras.initializers.Constant(6.0),
|
|
trainable=False,
|
|
)
|
|
|
|
def call(self, inputs):
|
|
x = tf.quantization.fake_quant_with_min_max_vars(
|
|
inputs, self.min_var, self.max_var
|
|
)
|
|
|
|
w_fq = tf.quantization.fake_quant_with_min_max_vars(
|
|
self.w, self.min_var, self.max_var
|
|
)
|
|
x = tf.matmul(x, w_fq)
|
|
|
|
x = tf.quantization.fake_quant_with_min_max_vars(
|
|
x, self.min_var, self.max_var
|
|
)
|
|
|
|
return x
|
|
|
|
return tf.keras.Sequential(QLinear(3, input_shape=(2,)))
|
|
|
|
@parameterized.named_parameters(
|
|
('_DefaultFLOAT32InputOutput', dtypes.float32),
|
|
('_INT8InputOutput', dtypes.int8),
|
|
('_UINT8InputOutput', dtypes.uint8),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testTrainingTimeQuantization(self, inference_input_output_type):
|
|
model = self._getTrainingTimeQuantizedModel()
|
|
|
|
float_converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
float_tflite_model = float_converter.convert()
|
|
self.assertIsNotNone(float_tflite_model)
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.inference_input_type = inference_input_output_type
|
|
quantized_converter.inference_output_type = inference_input_output_type
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(quantized_tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertAllEqual(
|
|
[metadata_fb.ModelOptimizationMode.QUANTIZATION_AWARE_TRAINING],
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual(
|
|
inference_input_output_type.as_numpy_dtype, input_details[0]['dtype']
|
|
)
|
|
output_details = interp.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertEqual(
|
|
inference_input_output_type.as_numpy_dtype, output_details[0]['dtype']
|
|
)
|
|
|
|
# Ensure that the quantized tflite model is smaller.
|
|
self.assertLess(len(quantized_tflite_model), len(float_tflite_model))
|
|
|
|
@test_util.run_v2_only
|
|
def testNewQuantizer(self):
|
|
"""Test the model quantized by the new converter."""
|
|
root, func, calibration_gen = self._getIntegerQuantizeModel()
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8
|
|
]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
|
|
# default quantizer
|
|
quantized_converter.experimental_new_quantizer = False
|
|
old_tflite = quantized_converter.convert()
|
|
|
|
# new quantizer
|
|
quantized_converter.experimental_new_quantizer = True
|
|
new_tflite = quantized_converter.convert()
|
|
|
|
for _ in range(5):
|
|
input_data = tf.constant(
|
|
np.random.uniform(-1, 1, size=(1, 5, 5, 3)).astype(np.float32)
|
|
)
|
|
old_value = self._evaluateTFLiteModel(old_tflite, [input_data])
|
|
new_value = self._evaluateTFLiteModel(new_tflite, [input_data])
|
|
self.assertAllClose(old_value, new_value, atol=1e-01)
|
|
|
|
@test_util.run_v2_only
|
|
def testGatherNDQI8(self):
|
|
"""Test gather_nd with quantized i8 parameters."""
|
|
|
|
class GatherNDQI8QDQ(tf.keras.Model):
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=(2, 2), dtype=tf.float32)]
|
|
)
|
|
def func(self, input_tensor):
|
|
x = tf.quantization.fake_quant_with_min_max_args(
|
|
input_tensor, -3.0, 3.0
|
|
)
|
|
x = tf.gather_nd(x, [[0, 0], [1, 1]])
|
|
return tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
|
|
|
|
# Build a QDQ model so that tfl.gather_nd will be converted to a QI8 version
|
|
# with the `_experimental_qdq_conversion_mode`` flag
|
|
root = GatherNDQI8QDQ()
|
|
concrete_func = root.func.get_concrete_function()
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
converter._experimental_qdq_conversion_mode = 'STATIC'
|
|
tflite_model = converter.convert()
|
|
|
|
np_data = np.array([[1, 2], [3, 4]], dtype=np.float32)
|
|
input_tensor = tf.constant(np_data, dtype=tf.int8)
|
|
expected_value = [1, 4]
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_tensor])
|
|
self.assertAllClose(expected_value, actual_value[0], atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testEmbeddings(self):
|
|
"""Test model with embeddings."""
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((20)), dtype=np.int32)
|
|
)
|
|
|
|
class EmbeddingModel(tf.keras.Model):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.shared_weights = self.add_weight(
|
|
'weights',
|
|
shape=(2000, 300),
|
|
dtype=tf.float32,
|
|
initializer=tf.random_normal_initializer(
|
|
mean=0.0, stddev=300 ** (-0.5)
|
|
),
|
|
)
|
|
|
|
@tf.function(input_signature=[tf.TensorSpec(shape=(20), dtype=tf.int32)])
|
|
def func(self, x):
|
|
return tf.gather(self.shared_weights, x)
|
|
|
|
# Building the model.
|
|
root = EmbeddingModel()
|
|
concrete_func = root.func.get_concrete_function()
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.func(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertAllClose(expected_value.numpy(), actual_value[0], atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testGraphDebugInfo(self):
|
|
"""Test a concrete function has debug info captured."""
|
|
root = autotrackable.AutoTrackable()
|
|
root.v1 = tf.Variable(3.0)
|
|
root.f = tf.function(lambda x: root.v1 * x)
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
converter.convert()
|
|
self._assertValidDebugInfo(converter._debug_info)
|
|
|
|
def _getIntegerQuantizationModelWithFlexOp(self):
|
|
np.random.seed(0)
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[3, 3, 3, 3, 3], dtype=tf.float32)]
|
|
)
|
|
def func(inp):
|
|
tanh = tf.math.tanh(inp)
|
|
# Flex delegate will merge the consecutive conv3d and erf ops into one
|
|
# Delegate node.
|
|
conv3d = tf.nn.conv3d(
|
|
tanh,
|
|
tf.ones([3, 3, 3, 3, 3]),
|
|
strides=[1, 1, 1, 1, 1],
|
|
padding='SAME',
|
|
)
|
|
erf = tf.math.erf(conv3d)
|
|
output = tf.math.tanh(erf)
|
|
return output
|
|
|
|
def calibration_gen():
|
|
for _ in range(5):
|
|
yield [
|
|
np.random.uniform(-1, 1, size=(3, 3, 3, 3, 3)).astype(np.float32)
|
|
]
|
|
|
|
root.f = func
|
|
return (root, root.f.get_concrete_function(), calibration_gen)
|
|
|
|
@parameterized.named_parameters(
|
|
('_Default', False, False, dtypes.float32),
|
|
('_INT8InputOutput', False, False, dtypes.int8),
|
|
('_UINT8InputOutput', False, False, dtypes.uint8),
|
|
('_INT16Quantize', False, True, dtypes.float32),
|
|
('_INT16Quantize_INT16InputOutput', False, True, dtypes.int16),
|
|
('_IntOnly', True, False, dtypes.float32),
|
|
('_IntOnly_INT8InputOutput', True, False, dtypes.int8),
|
|
('_IntOnly_UINT8InputOutput', True, False, dtypes.uint8),
|
|
('_IntOnly_INT16Quantize', True, True, dtypes.float32),
|
|
('_IntOnly_INT16Quantize_INT16InputOutput', True, True, dtypes.int16),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testIntegerQuantizationWithFlexOp(
|
|
self, is_int_only, is_int16_quantize, inference_input_output_type
|
|
):
|
|
root, func, calibration_gen = self._getIntegerQuantizationModelWithFlexOp()
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
if is_int_only:
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
else:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
else:
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
else:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
|
|
quantized_converter.inference_input_type = inference_input_output_type
|
|
quantized_converter.inference_output_type = inference_input_output_type
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(quantized_tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(metadata.options.enableSelectTfOps, True)
|
|
expected_opt_options = [metadata_fb.ModelOptimizationMode.PTQ_FULL_INTEGER]
|
|
if is_int16_quantize:
|
|
expected_opt_options = [metadata_fb.ModelOptimizationMode.PTQ_INT16]
|
|
self.assertAllEqual(
|
|
expected_opt_options, metadata.options.modelOptimizationModes
|
|
)
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual(
|
|
inference_input_output_type.as_numpy_dtype, input_details[0]['dtype']
|
|
)
|
|
output_details = interp.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertEqual(
|
|
inference_input_output_type.as_numpy_dtype, output_details[0]['dtype']
|
|
)
|
|
|
|
def _getIntegerQuantizationModelWithUnsupportedOps(self):
|
|
np.random.seed(0)
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[3], dtype=tf.float32),
|
|
tf.TensorSpec(shape=[3], dtype=tf.float32),
|
|
]
|
|
)
|
|
def func(a, b):
|
|
# ceil kernel does not support int8 nor int16 types neither.
|
|
left = tf.math.ceil(a)
|
|
right = tf.nn.tanh(b)
|
|
add = tf.math.add(left, right)
|
|
# ceil kernel does not support int8 nor int16 types neither.
|
|
output = tf.math.ceil(add)
|
|
return (output, right)
|
|
|
|
def calibration_gen():
|
|
for _ in range(5):
|
|
yield [
|
|
np.random.uniform(-1, 1, size=(3)).astype(np.float32),
|
|
np.random.uniform(-1, 1, size=(3)).astype(np.float32),
|
|
]
|
|
|
|
root.f = func
|
|
return (root, root.f.get_concrete_function(), calibration_gen)
|
|
|
|
@parameterized.named_parameters(
|
|
('_INT8InputOutput', False, False, dtypes.int8),
|
|
('_UINT8InputOutput', False, False, dtypes.uint8),
|
|
('_INT16Quantize_INT16InputOutput', False, True, dtypes.int16),
|
|
('_IntOnly_INT8InputOutput', True, False, dtypes.int8),
|
|
('_IntOnly_UINT8InputOutput', True, False, dtypes.uint8),
|
|
('_IntOnly_INT16Quantize_INT16InputOutput', True, True, dtypes.int16),
|
|
('_IntOnly_INT8InputOutputMlirQuant', True, False, dtypes.int8, True),
|
|
('_IntOnly_UINT8InputOutputMlirQuant', True, False, dtypes.uint8, True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testIntegerQuantizationWithUnsupportedOps(
|
|
self,
|
|
is_int_only,
|
|
is_int16_quantize,
|
|
inference_input_output_type,
|
|
enable_mlir_quantizer=False,
|
|
):
|
|
root, func, calib_gen = (
|
|
self._getIntegerQuantizationModelWithUnsupportedOps()
|
|
)
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.representative_dataset = calib_gen
|
|
if is_int_only:
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
]
|
|
else:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
]
|
|
else:
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
]
|
|
else:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS
|
|
]
|
|
|
|
quantized_converter.inference_input_type = inference_input_output_type
|
|
quantized_converter.inference_output_type = inference_input_output_type
|
|
quantized_converter.experimental_new_quantizer = enable_mlir_quantizer
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
expected_dtype = inference_input_output_type.as_numpy_dtype
|
|
# Allow float32 for fallback on non-quantizable op.
|
|
expected_ceil_dtype = (
|
|
expected_dtype if enable_mlir_quantizer else dtypes.float32
|
|
)
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 2)
|
|
self.assertEqual(input_details[0]['dtype'], expected_dtype)
|
|
self.assertEqual(input_details[1]['dtype'], expected_ceil_dtype)
|
|
output_details = interp.get_output_details()
|
|
self.assertLen(output_details, 2)
|
|
self.assertEqual(output_details[0]['dtype'], expected_dtype)
|
|
self.assertEqual(output_details[1]['dtype'], expected_ceil_dtype)
|
|
|
|
def _getIntegerQuantizationModelWithControlFlow(self):
|
|
def true_fn(x):
|
|
return x
|
|
|
|
def false_fn(x):
|
|
return x
|
|
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[1, 2], dtype=tf.float32),
|
|
tf.TensorSpec(shape=(), dtype=tf.bool),
|
|
]
|
|
)
|
|
def model(x, b):
|
|
x = x + x
|
|
x = tf.cond(b, true_fn=lambda: true_fn(x), false_fn=lambda: false_fn(x))
|
|
return x + x
|
|
|
|
def calibration_gen():
|
|
for _ in range(5):
|
|
yield [
|
|
np.random.uniform(
|
|
-1,
|
|
1,
|
|
size=(
|
|
1,
|
|
2,
|
|
),
|
|
).astype(np.float32),
|
|
tf.constant(True),
|
|
]
|
|
for _ in range(5):
|
|
yield [
|
|
np.random.uniform(
|
|
-1,
|
|
1,
|
|
size=(
|
|
1,
|
|
2,
|
|
),
|
|
).astype(np.float32),
|
|
tf.constant(False),
|
|
]
|
|
|
|
return (model, model.get_concrete_function(), calibration_gen)
|
|
|
|
@parameterized.named_parameters(
|
|
('_INT8InputOutput', False, False, dtypes.int8),
|
|
('_UINT8InputOutput', False, False, dtypes.uint8),
|
|
('_INT16Quantize_INT16InputOutput', False, True, dtypes.int16),
|
|
('_IntOnly_INT8InputOutput', True, False, dtypes.int8),
|
|
('_IntOnly_UINT8InputOutput', True, False, dtypes.uint8),
|
|
('_IntOnly_INT16Quantize_INT16InputOutput', True, True, dtypes.int16),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testIntegerQuantizationWithControlFlow(
|
|
self,
|
|
is_int_only,
|
|
is_int16_quantize,
|
|
inference_input_output_type,
|
|
enable_mlir_quantizer=False,
|
|
):
|
|
root, func, calib_gen = self._getIntegerQuantizationModelWithControlFlow()
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.representative_dataset = calib_gen
|
|
if is_int_only:
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
]
|
|
else:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
]
|
|
else:
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8,
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
]
|
|
else:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS
|
|
]
|
|
|
|
quantized_converter.inference_input_type = inference_input_output_type
|
|
quantized_converter.inference_output_type = inference_input_output_type
|
|
quantized_converter.experimental_new_quantizer = enable_mlir_quantizer
|
|
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
expected_dtype = inference_input_output_type.as_numpy_dtype
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 2)
|
|
self.assertEqual(input_details[0]['dtype'], expected_dtype)
|
|
self.assertEqual(input_details[1]['dtype'], dtypes.bool)
|
|
output_details = interp.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertEqual(output_details[0]['dtype'], expected_dtype)
|
|
|
|
@parameterized.named_parameters(
|
|
('_BlocklistedNoneWithLowering', None, None, True),
|
|
('_BlocklistedNoneWithoutLowering', None, None, False),
|
|
('_BlocklistedOpsWithLowering', {'CONV_2D'}, None, True),
|
|
('_BlocklistedOpsWithoutLowering', {'CONV_2D'}, None, False),
|
|
('_BlocklistedNodesWithLowering', None, {'PartitionedCall:0'}, True),
|
|
('_BlocklistedNodesWithoutLowering', None, {'Identity'}, False),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testNewQuantizerBlocklistingArgs(
|
|
self, denylisted_ops, denylisted_nodes, lower_to_saved_model
|
|
):
|
|
"""Test the model quantized by the new converter and denylisted options."""
|
|
root, func, calibration_gen = self._getIntegerQuantizeModel()
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8
|
|
]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.experimental_new_quantizer = True
|
|
quantized_converter._experimental_calibrate_only = True
|
|
quantized_converter.experimental_lower_to_saved_model = lower_to_saved_model
|
|
calibrated = quantized_converter.convert()
|
|
quantized_tflite_model = convert.mlir_quantize(
|
|
calibrated,
|
|
denylisted_ops=denylisted_ops,
|
|
denylisted_nodes=denylisted_nodes,
|
|
)
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
details = interp.get_tensor_details()
|
|
num_quantized_tensors = sum([
|
|
1
|
|
for detail in details
|
|
if len(detail['quantization_parameters']['scales'])
|
|
])
|
|
if denylisted_nodes or denylisted_ops:
|
|
self.assertEqual(num_quantized_tensors, 0)
|
|
return
|
|
self.assertEqual(num_quantized_tensors, 4) # quant, filter, bias, dequant
|
|
|
|
@parameterized.named_parameters(
|
|
('_SingleLayer', False),
|
|
('_WholeModel', True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testNewQuantizerNumericVerificationDebugMode(self, whole_model_verify):
|
|
"""Test the model quantized by the new converter with numeric verify ops."""
|
|
root, func, calibration_gen = self._getIntegerQuantizeModel()
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8
|
|
]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
|
|
# Create a TFLite model with new quantizer.
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.experimental_new_quantizer = True
|
|
production_tflite = quantized_converter.convert()
|
|
# Create a TFLite model with new quantizer and numeric verify ops.
|
|
quantized_converter._experimental_calibrate_only = True
|
|
calibrated = quantized_converter.convert()
|
|
debug_mode_tflite = convert.mlir_quantize(
|
|
calibrated,
|
|
enable_numeric_verify=True,
|
|
enable_whole_model_verify=whole_model_verify,
|
|
)
|
|
|
|
# Check if adding debug mode should output a different flatbuffer.
|
|
self.assertNotEqual(production_tflite, debug_mode_tflite)
|
|
|
|
# Check if newly added ops are numeric verify ops.
|
|
input_data = tf.constant(
|
|
np.random.uniform(-1, 1, size=(1, 5, 5, 3)).astype(np.float32)
|
|
)
|
|
|
|
def examine_tflite_model(tflite_content, input_data):
|
|
interp = interpreter.Interpreter(
|
|
model_content=tflite_content,
|
|
experimental_op_resolver_type=interpreter.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES,
|
|
)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
interp.set_tensor(input_details[0]['index'], input_data.numpy())
|
|
interp.invoke()
|
|
tensor_details = interp.get_tensor_details()
|
|
return {
|
|
details['name']: interp.get_tensor(details['index'])
|
|
for details in interp.get_tensor_details()
|
|
}, tensor_details
|
|
|
|
tflite_result, _ = examine_tflite_model(production_tflite, input_data)
|
|
debug_mode_tflite_result, debug_tensor_details = examine_tflite_model(
|
|
debug_mode_tflite, input_data
|
|
)
|
|
|
|
# MLIR-based quantizer should output flatbuffer model with `tfl.quantize`.
|
|
num_production_quantize_ops = len([
|
|
None
|
|
for output_tensor_name in tflite_result
|
|
if 'tfl.quantize' in output_tensor_name
|
|
])
|
|
self.assertEqual(num_production_quantize_ops, 1)
|
|
# MLIR-based quantizer should output flatbuffer model with `tfl.quantize`.
|
|
num_debug_quantize_ops = len([
|
|
None
|
|
for output_tensor_name in debug_mode_tflite_result
|
|
if 'tfl.quantize' in output_tensor_name
|
|
])
|
|
# Two numbers should be equal.
|
|
self.assertEqual(num_production_quantize_ops, num_debug_quantize_ops)
|
|
# DebugMode TFLite flatbuffer should have NumericVerifyOps more than zero.
|
|
# The name has the prefix "NumericVerify/{name}:{id}
|
|
# where {name} is the tensor name of the original quantized op's activation,
|
|
# and {id} is its tensor id.
|
|
num_debug_ops = 0
|
|
for output_tensor_name in debug_mode_tflite_result:
|
|
if 'NumericVerify' in output_tensor_name:
|
|
pos_end_prefix = len('NumericVerify/')
|
|
pos_colon = output_tensor_name.rfind(':')
|
|
self.assertEqual('NumericVerify/', output_tensor_name[:pos_end_prefix])
|
|
tensor_id = int(output_tensor_name[pos_colon + 1 :])
|
|
original_tensor_name = output_tensor_name[pos_end_prefix:pos_colon]
|
|
self.assertEqual(
|
|
original_tensor_name, debug_tensor_details[tensor_id]['name']
|
|
)
|
|
num_debug_ops += 1
|
|
self.assertEqual(num_debug_ops, 1)
|
|
# The number of debug ops should be equal to that of quantized ops.
|
|
self.assertEqual(num_debug_ops, num_debug_quantize_ops)
|
|
|
|
@parameterized.named_parameters(
|
|
('_PerChannelQuant', False, False),
|
|
('_PerChannelMlirQuant', False, True),
|
|
('_PerTensorQuant', True, False),
|
|
('_PerTensorMlirQuant', True, True),
|
|
('_PerChannelDynamicRange', False, False),
|
|
('_PerTensorDynamicRange', True, False),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testDisablePerChannelQuantization(
|
|
self,
|
|
disable_per_channel=False,
|
|
enable_mlir_quantizer=False,
|
|
):
|
|
k_conv_name = 'Conv2D'
|
|
# Dynamic range quant requires total num elements of filters > 1024.
|
|
k_num_filters = 38
|
|
root, func, calib_gen = self._getIntegerQuantizeModel(k_num_filters)
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.representative_dataset = calib_gen
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS
|
|
]
|
|
quantized_converter.experimental_new_quantizer = enable_mlir_quantizer
|
|
if disable_per_channel:
|
|
quantized_converter._experimental_disable_per_channel = (
|
|
disable_per_channel
|
|
)
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
detail = next((
|
|
d
|
|
for d in interp.get_tensor_details()
|
|
if d['name'].startswith(k_conv_name)
|
|
))
|
|
quant_params = detail['quantization_parameters']
|
|
expected_num_params = 1 if disable_per_channel else k_num_filters
|
|
self.assertLen(quant_params['scales'], expected_num_params)
|
|
if len(quant_params['zero_points']) != 1:
|
|
self.assertLen(quant_params['zero_points'], expected_num_params)
|
|
|
|
def _getIntegerQuantizeDenseModel(self, num_filters=32):
|
|
np.random.seed(0)
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[1, 16], dtype=tf.float32)]
|
|
)
|
|
def func(inp):
|
|
dense = tf.matmul(a=inp, b=tf.ones([16, num_filters]))
|
|
output = tf.nn.relu(dense, name='output')
|
|
return output
|
|
|
|
def calibration_gen():
|
|
for _ in range(5):
|
|
yield [np.random.uniform(-1, 1, size=(1, 16)).astype(np.float32)]
|
|
|
|
root.f = func
|
|
to_save = root.f.get_concrete_function()
|
|
return (root, to_save, calibration_gen)
|
|
|
|
@parameterized.named_parameters(
|
|
('_PerChannelQuant', False, False),
|
|
('_PerChannelMlirQuant', False, True),
|
|
('_PerTensorQuant', True, False),
|
|
('_PerTensorMlirQuant', True, True),
|
|
('_PerChannelDynamicRange', False, True, True),
|
|
('_PerTensorDynamicRange', True, True, True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testDisablePerChannelQuantizationForDenseLayers(
|
|
self,
|
|
disable_per_channel_for_dense=False,
|
|
enable_mlir_quantizer=False,
|
|
representative_dataset=False,
|
|
):
|
|
k_dense_name = 'MatMul'
|
|
# Dynamic range quant requires total num elements of filters > 1024.
|
|
k_num_filters = 64
|
|
root, func, calib_gen = self._getIntegerQuantizeDenseModel(k_num_filters)
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
if representative_dataset:
|
|
quantized_converter.representative_dataset = calib_gen
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS
|
|
]
|
|
quantized_converter.experimental_new_quantizer = enable_mlir_quantizer
|
|
if disable_per_channel_for_dense:
|
|
quantized_converter._experimental_disable_per_channel_quantization_for_dense_layers = (
|
|
disable_per_channel_for_dense
|
|
)
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
# Do not apply delegates as XNNPack converts per tensor to per channel.
|
|
interp = interpreter.Interpreter(
|
|
model_content=quantized_tflite_model,
|
|
experimental_op_resolver_type=interpreter.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES,
|
|
)
|
|
interp.allocate_tensors()
|
|
detail = next((
|
|
d
|
|
for d in interp.get_tensor_details()
|
|
if d['name'].startswith(k_dense_name)
|
|
))
|
|
quant_params = detail['quantization_parameters']
|
|
expected_num_params = 1 if disable_per_channel_for_dense else k_num_filters
|
|
self.assertLen(quant_params['scales'], expected_num_params)
|
|
if len(quant_params['zero_points']) != 1:
|
|
self.assertLen(quant_params['zero_points'], expected_num_params)
|
|
|
|
@parameterized.named_parameters(
|
|
('MlirQuantize', True), ('TocoQuantize', False)
|
|
)
|
|
@test_util.run_v2_only
|
|
def testQuantizeBiasOverflow(self, enable_mlir_quantizer):
|
|
"""Tests if the quantizer handles bias overflow by adjusting scales."""
|
|
input_data = np.array([[-1e-3, 1e-3]], dtype=np.float32)
|
|
|
|
def calibration_gen():
|
|
yield {'x': input_data}
|
|
|
|
root = self._getMatMulModelWithSmallWeights()
|
|
input_data = tf.constant([-1e-3, 1e-3], shape=(1, 2))
|
|
concrete_func = root.matmul.get_concrete_function(input_data)
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = calibration_gen
|
|
converter.experimental_new_quantizer = enable_mlir_quantizer
|
|
quantized_model = converter.convert()
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
interp.invoke()
|
|
output_details = interp.get_output_details()
|
|
output = interp.get_tensor(output_details[0]['index'])
|
|
# the inputs and weights are far smaller than the biases, so the final
|
|
# result should be equal to the biases.
|
|
self.assertAllClose(root.bias, output.flatten())
|
|
|
|
@test_util.run_v2_only
|
|
def testOpVersion(self):
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[5, 5], dtype=tf.float32)]
|
|
)
|
|
def custom_resize(image):
|
|
# Add "batch" and "channels" dimensions
|
|
image = image[tf.newaxis, ..., tf.newaxis]
|
|
# ResizeBilinear version 3.
|
|
resize1 = tf.compat.v1.image.resize_bilinear(
|
|
image, [2, 2], half_pixel_centers=True
|
|
)
|
|
# ResizeBilinear version 1.
|
|
resize2 = tf.compat.v1.image.resize_bilinear(image, [2, 2])
|
|
return resize1 + resize2
|
|
|
|
concrete_func = custom_resize.get_concrete_function()
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], custom_resize
|
|
)
|
|
tflite_model = converter.convert()
|
|
model_object = schema_fb.Model.GetRootAsModel(tflite_model, 0)
|
|
model = schema_fb.ModelT.InitFromObj(model_object)
|
|
|
|
for operator in model.operatorCodes:
|
|
if operator.builtinCode == schema_fb.BuiltinOperator.RESIZE_BILINEAR:
|
|
# half_pixel_centers is supported by ResizeBilinear version 3.
|
|
self.assertEqual(operator.version, 3)
|
|
break
|
|
|
|
@test_util.run_v2_only
|
|
def testForceSelectTFOps(self):
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
converter.target_spec.supported_ops = [lite.OpsSet.SELECT_TF_OPS]
|
|
tflite_model = converter.convert()
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(metadata.options.forceSelectTfOps, True)
|
|
|
|
# Check output value from converted model.
|
|
expected_value = root.f(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
def testExcludeConversionMetadata(self):
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
converter.exclude_conversion_metadata = True
|
|
tflite_model = converter.convert()
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(tflite_model)
|
|
self.assertIsNone(metadata)
|
|
|
|
def testConversionMetadataForDynamicRange(self):
|
|
func, _ = self._getCeilModel()
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func.get_concrete_function()]
|
|
)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_model = converter.convert()
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(quantized_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertAllEqual(
|
|
[metadata_fb.ModelOptimizationMode.PTQ_DYNAMIC_RANGE],
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
|
|
def testConversionMetadataForFloat16(self):
|
|
root, func, calibration_gen = self._getIntegerQuantizeModel()
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions([func], root)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = calibration_gen
|
|
converter.target_spec.supported_types = [dtypes.float16]
|
|
quantized_model = converter.convert()
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(quantized_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertAllEqual(
|
|
[metadata_fb.ModelOptimizationMode.PTQ_FLOAT16],
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
|
|
def testSerializeDebugMetadata(self):
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
converter.serialize_debug_metadata = True
|
|
tflite_model = flatbuffer_utils.convert_bytearray_to_object(
|
|
converter.convert()
|
|
)
|
|
|
|
# Check the debug metadata.
|
|
metadata_names = [m.name for m in tflite_model.metadata]
|
|
self.assertIn(b'debug_metadata', metadata_names)
|
|
|
|
|
|
class FromSavedModelTest(lite_v2_test_util.ModelTest):
|
|
|
|
def _createV1SavedModel(self, shape):
|
|
"""Create a simple SavedModel."""
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_savedmodel')
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor_1 = tf.compat.v1.placeholder(
|
|
shape=shape, dtype=tf.float32, name='inputB'
|
|
)
|
|
in_tensor_2 = tf.compat.v1.placeholder(
|
|
shape=shape, dtype=tf.float32, name='inputA'
|
|
)
|
|
variable_node = tf.Variable(1.0, name='variable_node')
|
|
out_tensor = in_tensor_1 + in_tensor_2 * variable_node
|
|
inputs = {'x': in_tensor_1, 'y': in_tensor_2}
|
|
outputs = {'z': out_tensor}
|
|
sess.run(tf.compat.v1.variables_initializer([variable_node]))
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
def _createV2QATSavedModel(self, shape):
|
|
"""Create a simple QAT SavedModel in TF 2."""
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
input_name = 'input'
|
|
output_name = 'scores'
|
|
|
|
class _FakeQuantArgsLayer(tf.keras.layers.Layer):
|
|
"""A fake quantization layer with fake_quant_with_min_max_args.
|
|
|
|
Keras 3 requires wrapping the tf function inside Keras layer.
|
|
"""
|
|
|
|
def call(self, x):
|
|
return tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
|
|
|
|
input_tensor = tf.keras.layers.Input((32, 32, 128), name=input_name)
|
|
x = _FakeQuantArgsLayer()(input_tensor)
|
|
x = tf.keras.layers.Conv2D(1, (3, 3))(x)
|
|
x = _FakeQuantArgsLayer()(x)
|
|
scores = tf.keras.layers.Reshape((-1,), name=output_name)(x)
|
|
model = tf.keras.Model(input_tensor, scores)
|
|
model.save(saved_model_dir)
|
|
return saved_model_dir, input_name, output_name
|
|
|
|
@test_util.run_v2_only
|
|
def testStableHloQuantizerSupportsOnlyStaticRangePtq(self):
|
|
"""Tests that StableHLO Quantizer supports only static-range PTQ."""
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = tf.function(lambda x: 2.0 * x)
|
|
to_save = root.f.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, to_save)
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
converter.experimental_use_stablehlo_quantizer = True
|
|
with self.assertRaisesRegex(
|
|
ValueError, 'only supports static-range and weight-only PTQ'
|
|
):
|
|
converter.convert()
|
|
|
|
@test_util.run_v2_only
|
|
def testStableHloQuantizerNoOpForTfSavedModel(self):
|
|
"""Tests that StableHLO Quantizer does not run for TF SavedModel."""
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = tf.function(lambda x: 2.0 * x)
|
|
to_save = root.f.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, to_save)
|
|
|
|
def _representative_data_gen():
|
|
return [{'x': np.ones(shape=(1,), dtype=np.float32)}]
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
# Set the flags to enable StableHLO Quantizer.
|
|
converter.experimental_use_stablehlo_quantizer = True
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = _representative_data_gen
|
|
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Test that no tensor is quantized.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
all_tensor_details = interp.get_tensor_details()
|
|
for tensor_detail in all_tensor_details:
|
|
self.assertIn('dtype', tensor_detail)
|
|
self.assertEqual(tensor_detail['dtype'], np.float32)
|
|
|
|
@test_util.run_v2_only
|
|
def testV1SimpleModel(self):
|
|
"""Test a SavedModel."""
|
|
with tf.Graph().as_default():
|
|
saved_model_dir = self._createV1SavedModel(shape=[1, 16, 16, 3])
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
interp.allocate_tensors()
|
|
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 2)
|
|
self.assertStartsWith(input_details[0]['name'], 'inputA')
|
|
self.assertEqual(np.float32, input_details[0]['dtype'])
|
|
self.assertAllEqual([1, 16, 16, 3], input_details[0]['shape'])
|
|
self.assertEqual((0.0, 0.0), input_details[0]['quantization'])
|
|
|
|
self.assertStartsWith(
|
|
input_details[1]['name'],
|
|
'inputB',
|
|
)
|
|
self.assertEqual(np.float32, input_details[1]['dtype'])
|
|
self.assertTrue([1, 16, 16, 3], input_details[1]['shape'])
|
|
self.assertEqual((0.0, 0.0), input_details[1]['quantization'])
|
|
|
|
output_details = interp.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertStartsWith(output_details[0]['name'], 'add')
|
|
self.assertEqual(np.float32, output_details[0]['dtype'])
|
|
self.assertTrue([1, 16, 16, 3], output_details[0]['shape'])
|
|
self.assertEqual((0.0, 0.0), output_details[0]['quantization'])
|
|
|
|
@parameterized.named_parameters(
|
|
('Default', False),
|
|
('UnfoldLargeConstant', True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testUnfoldLargeConstant(self, unfold_large_constant):
|
|
"""Test unfolding large splat constant in a TF Lite model."""
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_savedmodel')
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1000, 1000], dtype=tf.float32, name='input'
|
|
)
|
|
constant = tf.constant(value=1, dtype=tf.float32, shape=[1000, 1000])
|
|
out_tensor = in_tensor + constant
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'y': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter._experimental_unfold_large_splat_constant = unfold_large_constant
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
model = util._convert_model_from_bytearray_to_object(tflite_model)
|
|
if unfold_large_constant:
|
|
self.assertEqual(
|
|
model.operatorCodes[0].builtinCode, schema_fb.BuiltinOperator.FILL
|
|
)
|
|
self.assertEqual(
|
|
model.operatorCodes[1].builtinCode, schema_fb.BuiltinOperator.ADD
|
|
)
|
|
else:
|
|
self.assertEqual(
|
|
model.operatorCodes[0].builtinCode, schema_fb.BuiltinOperator.ADD
|
|
)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
interp.allocate_tensors()
|
|
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual('input:0', input_details[0]['name'])
|
|
self.assertEqual(np.float32, input_details[0]['dtype'])
|
|
self.assertAllEqual([1000, 1000], input_details[0]['shape'])
|
|
self.assertEqual((0.0, 0.0), input_details[0]['quantization'])
|
|
|
|
output_details = interp.get_output_details()
|
|
self.assertEqual('add:0', output_details[0]['name'])
|
|
self.assertEqual(np.float32, output_details[0]['dtype'])
|
|
self.assertAllEqual([1000, 1000], output_details[0]['shape'])
|
|
self.assertEqual((0.0, 0.0), output_details[0]['quantization'])
|
|
|
|
interp.set_tensor(
|
|
input_details[0]['index'], np.ones(shape=[1000, 1000], dtype=np.float32)
|
|
)
|
|
interp.invoke()
|
|
self.assertAllEqual(
|
|
np.full(shape=[1000, 1000], fill_value=2.0, dtype=np.float32),
|
|
interp.get_tensor(output_details[0]['index']),
|
|
)
|
|
|
|
@test_util.run_v2_only
|
|
def testPreserveAssert(self):
|
|
"""Test preserving AssertOp in a TF Lite model."""
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_savedmodel')
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[10, 10], dtype=tf.float32, name='input'
|
|
)
|
|
constant = tf.constant(value=1, dtype=tf.float32, shape=[10, 10])
|
|
assert_op = tf.Assert(tf.less_equal(in_tensor, constant), [in_tensor])
|
|
with tf.control_dependencies([assert_op]):
|
|
out_tensor = in_tensor + constant
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'y': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
converter._experimental_preserve_assert_op = True
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
model = util._convert_model_from_bytearray_to_object(tflite_model)
|
|
has_assert = False
|
|
for op_code in model.operatorCodes:
|
|
if op_code.customCode == b'FlexAssert':
|
|
has_assert = True
|
|
break
|
|
self.assertTrue(has_assert)
|
|
|
|
@test_util.run_v2_only
|
|
def testTF1HubFormattedModel(self):
|
|
"""Test a TF1 hub formatted model."""
|
|
saved_model_dir = self._createV1SavedModel(shape=[1, 16, 16, 3])
|
|
|
|
# TF1 hub model is based on V1 saved model and they omit the saved model
|
|
# schema version setting.
|
|
saved_model_proto = loader_impl.parse_saved_model(saved_model_dir)
|
|
saved_model_proto.saved_model_schema_version = 0
|
|
|
|
saved_model_pb_file_path = os.path.join(saved_model_dir, 'saved_model.pb')
|
|
with file_io.FileIO(saved_model_pb_file_path, 'wb') as writer:
|
|
writer.write(saved_model_proto.SerializeToString())
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
def _createV1ModelWithHashTableInitializer(self):
|
|
# Create a v1 saved model with hash table initializers.
|
|
tf.compat.v1.disable_eager_execution()
|
|
saved_model_dir = os.path.join(
|
|
self.get_temp_dir(), 'savedmodel_with_hashtable'
|
|
)
|
|
|
|
table_initializer = tf.lookup.KeyValueTensorInitializer(
|
|
keys=['a', 'b', 'c', 'd'],
|
|
values=[1, 2, 3, 4],
|
|
key_dtype=tf.string,
|
|
value_dtype=tf.int64,
|
|
)
|
|
table = tf.lookup.StaticHashTable(
|
|
table_initializer, default_value=tf.constant(-1, dtype=tf.int64)
|
|
)
|
|
|
|
x = tf.compat.v1.placeholder(tf.string, shape=(), name='input')
|
|
y = table.lookup(x)
|
|
|
|
tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(x)
|
|
tensor_info_y = tf.compat.v1.saved_model.utils.build_tensor_info(y)
|
|
|
|
signature_def_map, init_op, assets_collection = (
|
|
{
|
|
'serving_default': tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
|
|
inputs={'x': tensor_info_x},
|
|
outputs={'y': tensor_info_y},
|
|
method_name='some_function',
|
|
)
|
|
},
|
|
tf.compat.v1.tables_initializer(),
|
|
None,
|
|
)
|
|
|
|
sess = tf.compat.v1.Session()
|
|
sess.run(tf.compat.v1.initializers.global_variables())
|
|
|
|
builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(
|
|
saved_model_dir
|
|
)
|
|
builder.add_meta_graph_and_variables(
|
|
sess,
|
|
[tf.compat.v1.saved_model.tag_constants.SERVING],
|
|
signature_def_map,
|
|
main_op=init_op,
|
|
assets_collection=assets_collection,
|
|
strip_default_attrs=True,
|
|
)
|
|
builder.save()
|
|
|
|
# Restore TF v2 behavior.
|
|
tf.compat.v1.reset_default_graph()
|
|
tf.compat.v1.enable_eager_execution()
|
|
return saved_model_dir
|
|
|
|
@test_util.run_v2_only
|
|
def testModelWithHashTableInitializer(self):
|
|
"""Test a model with saved_model's session initializer for hash tables."""
|
|
saved_model_dir = self._createV1ModelWithHashTableInitializer()
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
input_data = np.array(['a', 'b', 'c', 'z'], dtype=np.bytes_)
|
|
interp.resize_tensor_input(input_details[0]['index'], [4], strict=False)
|
|
interp.allocate_tensors()
|
|
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
# Invoke multiple times to ensure the initializer graph runs only once.
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual([1, 2, 3, -1], list(actual_value))
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual([1, 2, 3, -1], list(actual_value))
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual([1, 2, 3, -1], list(actual_value))
|
|
|
|
def _createV1ModelWithMutableHashTable(self):
|
|
# Create a v1 saved model with mutable hash table.
|
|
tf.compat.v1.disable_eager_execution()
|
|
saved_model_dir = os.path.join(
|
|
self.get_temp_dir(), 'savedmodel_with_mutable_hashtable'
|
|
)
|
|
|
|
table = tf.raw_ops.MutableHashTableV2(
|
|
key_dtype=tf.string, value_dtype=tf.int64
|
|
)
|
|
x = tf.compat.v1.placeholder(tf.string, shape=(), name='input')
|
|
keys = tf.constant(['a', 'b'], tf.string)
|
|
values = tf.constant([1, 5], tf.int64)
|
|
default_value = tf.constant(-1, tf.int64)
|
|
insert_call = tf.raw_ops.LookupTableInsertV2(
|
|
table_handle=table, keys=keys, values=values
|
|
)
|
|
with tf.control_dependencies([insert_call]):
|
|
y = tf.raw_ops.LookupTableFindV2(
|
|
table_handle=table, keys=x, default_value=default_value
|
|
)
|
|
|
|
tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(x)
|
|
tensor_info_y = tf.compat.v1.saved_model.utils.build_tensor_info(y)
|
|
|
|
signature_def_map, init_op, assets_collection = (
|
|
{
|
|
'serving_default': tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
|
|
inputs={'x': tensor_info_x},
|
|
outputs={'y': tensor_info_y},
|
|
method_name='some_function',
|
|
)
|
|
},
|
|
tf.compat.v1.tables_initializer(),
|
|
None,
|
|
)
|
|
|
|
sess = tf.compat.v1.Session()
|
|
|
|
builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(
|
|
saved_model_dir
|
|
)
|
|
builder.add_meta_graph_and_variables(
|
|
sess,
|
|
[tf.compat.v1.saved_model.tag_constants.SERVING],
|
|
signature_def_map,
|
|
main_op=init_op,
|
|
assets_collection=assets_collection,
|
|
strip_default_attrs=True,
|
|
)
|
|
builder.save()
|
|
|
|
# Restore TF v2 behavior.
|
|
tf.compat.v1.reset_default_graph()
|
|
tf.compat.v1.enable_eager_execution()
|
|
return saved_model_dir
|
|
|
|
@test_util.run_v2_only
|
|
def testModelWithMutableHashTable(self):
|
|
"""Test a model with saved_model's session initializer for hash tables."""
|
|
saved_model_dir = self._createV1ModelWithMutableHashTable()
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
input_data = np.array(['a', 'b', 'c'], dtype=np.bytes_)
|
|
interp.resize_tensor_input(input_details[0]['index'], [3], strict=False)
|
|
interp.allocate_tensors()
|
|
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual([1, 5, -1], list(actual_value))
|
|
|
|
@test_util.run_v2_only
|
|
def testReduceSumWithInt16Quant(self):
|
|
"""Test a model with quantized int16 reduce sum op."""
|
|
inp = tf.keras.Input([3, 3], 3, name='x')
|
|
m = tf.keras.Model(inp, tf.reduce_sum(inp, axis=-1))
|
|
|
|
converter = lite.TFLiteConverterV2.from_keras_model(m)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
|
|
]
|
|
converter.inference_input_type = tf.int16
|
|
converter.inference_output_type = tf.int16
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
inputs = {
|
|
i.name: np.random.normal(size=i.shape).astype(np.float32)
|
|
for i in m.inputs
|
|
}
|
|
converter.representative_dataset = lambda: [inputs]
|
|
content = converter.convert()
|
|
|
|
interp = interpreter.Interpreter(model_content=content)
|
|
runner = interp.get_signature_runner('serving_default')
|
|
y = runner(x=np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]]).astype(np.int16))
|
|
self.assertEqual([3, 6, 9], list(list(y.values())[0]))
|
|
|
|
@test_util.run_v2_only
|
|
def testConstModel(self):
|
|
"""Test a basic model with functions to make sure functions are inlined."""
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = tf.function(lambda x: 2.0 * x)
|
|
to_save = root.f.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, to_save)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.f(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testVariableModel(self):
|
|
"""Test a basic model with Variables with saving/loading the SavedModel."""
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
to_save = root.f.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, to_save)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
tflite_model = converter.convert()
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(
|
|
metadata.environment.modelType, metadata_fb.ModelType.TF_SAVED_MODEL
|
|
)
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.f(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
@parameterized.named_parameters(
|
|
('EnableResourceVariables', True), ('DisableResourceVariables', False)
|
|
)
|
|
@test_util.run_v2_only
|
|
def testNativeVariablesModel(self, enable_resource_variables):
|
|
"""Test a basic model with Variables with saving/loading the SavedModel."""
|
|
root = self._getSimpleModelWithVariables()
|
|
input_data = tf.constant(1.0, shape=[1, 10])
|
|
to_save = root.assign_add.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, to_save)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
converter.experimental_enable_resource_variables = enable_resource_variables
|
|
|
|
# TODO(b/355497070): Remove this check once the
|
|
# CreateFreezeGlobalTensorsPass is migrated to the new TFL::Pass
|
|
# in the converter.
|
|
|
|
# if not enable_resource_variables:
|
|
# with self.assertRaises(convert.ConverterError) as error:
|
|
# tflite_model = converter.convert()
|
|
# self.assertIn(
|
|
# 'is not immutable, try removing mutable variables in your model
|
|
# ' since mutable variables are currently not supported through this'
|
|
# ' converter',
|
|
# str(error.exception),
|
|
# )
|
|
# return
|
|
|
|
# Enable resource variables.
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.assign_add(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
for tf_result, tflite_result in zip(expected_value, actual_value[0]):
|
|
self.assertAllClose(tf_result, tflite_result, atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testSignatures(self):
|
|
"""Test values for `signature_keys` argument."""
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
to_save = root.f.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, to_save)
|
|
|
|
# Convert model with invalid `signature_keys`.
|
|
with self.assertRaises(ValueError) as error:
|
|
_ = lite.TFLiteConverterV2.from_saved_model(
|
|
save_dir, signature_keys=['INVALID']
|
|
)
|
|
self.assertIn("Invalid signature key 'INVALID'", str(error.exception))
|
|
|
|
# Convert model with empty `signature_keys`.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(
|
|
save_dir, signature_keys=[]
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.f(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testSignatureDefsWithFullIntegerQuantization(self):
|
|
# SETUP
|
|
# 1. Define input shapes
|
|
tf_input_shape = (32, 32, 128)
|
|
tflite_input_shape = (1,) + tf_input_shape
|
|
# 2. Define model
|
|
tf_saved_model_dir, input_name, output_name = self._createV2QATSavedModel(
|
|
tf_input_shape
|
|
)
|
|
|
|
# MODEL 1: TFLite (float) model
|
|
# 1. Create TFLite model
|
|
converter = lite.TFLiteConverterV2.from_saved_model(tf_saved_model_dir)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
tflite_model = converter.convert()
|
|
# 2. Initialize the Interpreter
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()[0]
|
|
interp.resize_tensor_input(input_details['index'], tflite_input_shape)
|
|
interp.allocate_tensors()
|
|
# 3. (Skip) Verify that signature def input/output tensors are in the model.
|
|
# 4. Evaluate the model
|
|
input_data = np.random.random(tflite_input_shape).astype(np.float32)
|
|
result = self._evaluateTFLiteModelUsingSignatureDef(
|
|
tflite_model, 'serving_default', {input_name: input_data}
|
|
)[output_name]
|
|
|
|
# MODEL 2: TFLite (full integer quantized) model
|
|
# 1. Create TFLite model
|
|
converter = lite.TFLiteConverterV2.from_saved_model(tf_saved_model_dir)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.inference_input_type = tf.int8
|
|
converter.inference_output_type = tf.int8
|
|
tflite_model_quant = converter.convert()
|
|
# 2. Initialize the Interpreter
|
|
interp = interpreter.Interpreter(model_content=tflite_model_quant)
|
|
input_details = interp.get_input_details()[0]
|
|
output_details = interp.get_output_details()[0]
|
|
interp.resize_tensor_input(input_details['index'], tflite_input_shape)
|
|
interp.allocate_tensors()
|
|
# 3. Verify that signature def input/output tensors are in the model.
|
|
all_indices = {item['index'] for item in interp.get_tensor_details()}
|
|
signature_list = interp._get_full_signature_list()['serving_default']
|
|
input_tensor_indices = set(signature_list['inputs'].values())
|
|
assert input_tensor_indices.issubset(all_indices)
|
|
output_tensor_indices = set(signature_list['outputs'].values())
|
|
assert output_tensor_indices.issubset(all_indices)
|
|
|
|
# 4. Evaluate the model
|
|
input_data = np.random.random(tflite_input_shape)
|
|
input_scale, input_zero_point = input_details['quantization']
|
|
if (input_scale, input_zero_point) != (0.0, 0):
|
|
input_data = input_data / input_scale + input_zero_point
|
|
input_data = input_data.astype(input_details['dtype'])
|
|
result_quant = self._evaluateTFLiteModelUsingSignatureDef(
|
|
tflite_model_quant, 'serving_default', {input_name: input_data}
|
|
)[output_name]
|
|
output_scale, output_zero_point = output_details['quantization']
|
|
if (output_scale, output_zero_point) != (0.0, 0):
|
|
result_quant = result_quant.astype(np.float32)
|
|
result_quant = (result_quant - output_zero_point) * output_scale
|
|
|
|
# COMPARE: Validate that results from both models are approx. the same.
|
|
root_mean_squared = np.sqrt(np.mean((result - result_quant) ** 2))
|
|
assert root_mean_squared < 1.0
|
|
|
|
@test_util.run_v2_only
|
|
def testSignatureDefs(self):
|
|
"""Test converting SignatureDef is correct and uses SignatureDef API."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data_0 = tf.constant(1.0, shape=[1])
|
|
input_data_1 = tf.constant(3.0, shape=[1])
|
|
mul_add_func = root.mul_add.get_concrete_function(
|
|
input_data_1, input_data_0
|
|
)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, {'mul_add': mul_add_func})
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(
|
|
save_dir, signature_keys=['mul_add']
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.mul_add(input_data_1, input_data_0)
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
results = self._evaluateTFLiteModelUsingSignatureDef(
|
|
tflite_model, 'mul_add', {'y': input_data_0, 'x': input_data_1}
|
|
)
|
|
self.assertEqual(list(results.keys()), ['output_0'])
|
|
self.assertEqual(expected_value.numpy(), results['output_0'])
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 1)
|
|
self.assertEqual(list(signature_defs.keys()), ['mul_add'])
|
|
self.assertLen(signature_defs.values(), 1)
|
|
self.assertEqual(
|
|
list(signature_defs['mul_add'].keys()), ['inputs', 'outputs']
|
|
)
|
|
self.assertCountEqual(signature_defs['mul_add']['inputs'], ['x', 'y'])
|
|
self.assertEqual(list(signature_defs['mul_add']['outputs']), ['output_0'])
|
|
|
|
@test_util.run_v2_only
|
|
def testSignatureDefsWithDefaultValue(self):
|
|
"""Test converting SignatureDef is correct and uses SignatureDef API.
|
|
|
|
This test uses None as signature_key to test default behavior.
|
|
"""
|
|
root = self._getMultiFunctionModel()
|
|
input_data_0 = tf.constant(1.0, shape=[1])
|
|
input_data_1 = tf.constant(3.0, shape=[1])
|
|
mul_add_func = root.mul_add.get_concrete_function(
|
|
input_data_1, input_data_0
|
|
)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, {'mul_add': mul_add_func})
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(
|
|
save_dir, signature_keys=['mul_add']
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.mul_add(input_data_1, input_data_0)
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
results = self._evaluateTFLiteModelUsingSignatureDef(
|
|
tflite_model, None, {'y': input_data_0, 'x': input_data_1}
|
|
)
|
|
self.assertEqual(list(results.keys()), ['output_0'])
|
|
self.assertEqual(expected_value.numpy(), results['output_0'])
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 1)
|
|
self.assertEqual(list(signature_defs.keys()), ['mul_add'])
|
|
self.assertLen(signature_defs.values(), 1)
|
|
self.assertEqual(
|
|
list(signature_defs['mul_add'].keys()), ['inputs', 'outputs']
|
|
)
|
|
self.assertCountEqual(signature_defs['mul_add']['inputs'], ['x', 'y'])
|
|
self.assertEqual(list(signature_defs['mul_add']['outputs']), ['output_0'])
|
|
|
|
@test_util.run_v2_only
|
|
def testSignatureDefsQuantizedModel(self):
|
|
"""Test converting SignatureDef on quantized model."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data_0 = tf.constant(1.0, shape=[1])
|
|
input_data_1 = tf.constant(3.0, shape=[1])
|
|
mul_add_func = root.mul_add.get_concrete_function(
|
|
input_data_1, input_data_0
|
|
)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, {'mul_add': mul_add_func})
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(
|
|
save_dir, signature_keys=['mul_add']
|
|
)
|
|
|
|
def representative_dataset_gen():
|
|
for _ in range(2):
|
|
yield {
|
|
'x': np.random.uniform(low=0, high=1, size=(1, 1)).astype(
|
|
np.float32
|
|
),
|
|
'y': np.random.uniform(low=0, high=1, size=(1, 1)).astype(
|
|
np.float32
|
|
),
|
|
}
|
|
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = representative_dataset_gen
|
|
converter.target_spec.supported_ops = [lite.OpsSet.TFLITE_BUILTINS_INT8]
|
|
tflite_model = converter.convert()
|
|
|
|
# Check signatures are valid from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 1)
|
|
self.assertEqual(list(signature_defs.keys()), ['mul_add'])
|
|
self.assertLen(signature_defs.values(), 1)
|
|
self.assertEqual(
|
|
list(signature_defs['mul_add'].keys()), ['inputs', 'outputs']
|
|
)
|
|
self.assertCountEqual(signature_defs['mul_add']['inputs'], ['x', 'y'])
|
|
self.assertEqual(list(signature_defs['mul_add']['outputs']), ['output_0'])
|
|
|
|
@test_util.run_v2_only
|
|
def testMultipleFunctionModel(self):
|
|
"""Convert multiple functions in a multi-functional model."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
add_func = root.add.get_concrete_function(input_data)
|
|
sub_func = root.sub.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, {'add': add_func, 'sub': sub_func})
|
|
|
|
# Try converting multiple functions.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 2)
|
|
self.assertEqual(list(signature_defs.keys()), ['add', 'sub'])
|
|
self.assertLen(signature_defs.values(), 2)
|
|
self.assertEqual(list(signature_defs['add'].keys()), ['inputs', 'outputs'])
|
|
self.assertCountEqual(signature_defs['add']['inputs'], ['x'])
|
|
self.assertEqual(list(signature_defs['add']['outputs']), ['output_0'])
|
|
self.assertEqual(list(signature_defs['sub'].keys()), ['inputs', 'outputs'])
|
|
self.assertCountEqual(signature_defs['sub']['inputs'], ['x'])
|
|
self.assertEqual(list(signature_defs['sub']['outputs']), ['output_0'])
|
|
|
|
# Verify the Signature runner executions.
|
|
add_signature_runner = interp.get_signature_runner('add')
|
|
add_output = add_signature_runner(x=input_data)
|
|
self.assertEqual(add_output['output_0'], 3)
|
|
|
|
sub_signature_runner = interp.get_signature_runner('sub')
|
|
sub_output = sub_signature_runner(x=input_data)
|
|
self.assertEqual(sub_output['output_0'], -2)
|
|
|
|
@parameterized.named_parameters(
|
|
('_Default', False, False, dtypes.float32, False),
|
|
('_DefaultMlirQuant', False, False, dtypes.float32, True),
|
|
('_INT8InputOutput', False, False, dtypes.int8),
|
|
('_UINT8InputOutput', False, False, dtypes.uint8),
|
|
('_INT16Quantize_INT16InputOutput', False, True, dtypes.int16),
|
|
('_IntOnly_INT8InputOutput', True, False, dtypes.int8),
|
|
('_IntOnly_UINT8InputOutput', True, False, dtypes.uint8),
|
|
('_IntOnly_INT16Quantize_INT16InputOutput', True, True, dtypes.int16),
|
|
('_IntOnly_INT8InputOutputMlirQuant', True, False, dtypes.int8, True),
|
|
('_IntOnly_UINT8InputOutputMlirQuant', True, False, dtypes.uint8, True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testMultipleFunctionQuantizedModel(
|
|
self,
|
|
is_int_only,
|
|
is_int16_quantize,
|
|
inference_input_output_type,
|
|
enable_mlir_quantizer=False,
|
|
):
|
|
"""Convert multiple functions in a multi-functional model."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
add_func = root.add.get_concrete_function(input_data)
|
|
sub_func = root.sub.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, {'add': add_func, 'sub': sub_func})
|
|
|
|
# Try converting multiple functions.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
|
|
def representative_dataset_gen():
|
|
for _ in range(2):
|
|
yield (
|
|
'add',
|
|
{
|
|
'x': np.random.uniform(low=0, high=1, size=(1,)).astype(
|
|
np.float32
|
|
),
|
|
},
|
|
)
|
|
for _ in range(2):
|
|
yield (
|
|
'sub',
|
|
{
|
|
'x': np.random.uniform(low=0, high=1, size=(1,)).astype(
|
|
np.float32
|
|
),
|
|
},
|
|
)
|
|
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = representative_dataset_gen
|
|
if is_int_only:
|
|
if is_int16_quantize:
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
|
|
]
|
|
else:
|
|
converter.target_spec.supported_ops = [lite.OpsSet.TFLITE_BUILTINS_INT8]
|
|
else:
|
|
if is_int16_quantize:
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
|
|
]
|
|
else:
|
|
converter.target_spec.supported_ops = [lite.OpsSet.TFLITE_BUILTINS]
|
|
converter.inference_input_type = inference_input_output_type
|
|
converter.inference_output_type = inference_input_output_type
|
|
converter.experimental_new_quantizer = enable_mlir_quantizer
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 2)
|
|
self.assertEqual(list(signature_defs.keys()), ['add', 'sub'])
|
|
self.assertLen(signature_defs.values(), 2)
|
|
self.assertEqual(list(signature_defs['add'].keys()), ['inputs', 'outputs'])
|
|
self.assertCountEqual(signature_defs['add']['inputs'], ['x'])
|
|
self.assertEqual(list(signature_defs['add']['outputs']), ['output_0'])
|
|
self.assertEqual(list(signature_defs['sub'].keys()), ['inputs', 'outputs'])
|
|
self.assertCountEqual(signature_defs['sub']['inputs'], ['x'])
|
|
self.assertEqual(list(signature_defs['sub']['outputs']), ['output_0'])
|
|
|
|
# Verify the Signature runner executions.
|
|
input_data = tf.constant(
|
|
np.random.uniform(-1, 1, size=(1,)).astype(
|
|
inference_input_output_type.as_numpy_dtype
|
|
)
|
|
)
|
|
add_signature_runner = interp.get_signature_runner('add')
|
|
add_output = add_signature_runner(x=input_data)
|
|
self.assertIsNotNone(add_output['output_0'])
|
|
input_details = add_signature_runner.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertStartsWith(input_details['x']['name'], 'add_x:0')
|
|
self.assertEqual(
|
|
inference_input_output_type.as_numpy_dtype, input_details['x']['dtype']
|
|
)
|
|
self.assertTrue(([1] == input_details['x']['shape']).all())
|
|
if inference_input_output_type == dtypes.float32:
|
|
self.assertEqual((0.0, 0), input_details['x']['quantization'])
|
|
|
|
sub_signature_runner = interp.get_signature_runner('sub')
|
|
sub_output = sub_signature_runner(x=input_data)
|
|
self.assertIsNotNone(sub_output['output_0'])
|
|
output_details = sub_signature_runner.get_output_details()
|
|
self.assertLen(output_details, 1)
|
|
self.assertStartsWith(
|
|
output_details['output_0']['name'], 'StatefulPartitionedCall_1:0'
|
|
)
|
|
self.assertEqual(
|
|
inference_input_output_type.as_numpy_dtype,
|
|
output_details['output_0']['dtype'],
|
|
)
|
|
self.assertTrue(([1] == output_details['output_0']['shape']).all())
|
|
if inference_input_output_type == dtypes.float32:
|
|
self.assertEqual((0.0, 0), output_details['output_0']['quantization'])
|
|
|
|
@test_util.run_v2_only
|
|
def testMultipleFunctionModelWithSharedWeight(self):
|
|
"""Convert multiple functions with the shared weight."""
|
|
root = self._getMultiFunctionModelWithSharedWeight()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
add_func = root.add.get_concrete_function(input_data)
|
|
sub_func = root.sub.get_concrete_function(input_data)
|
|
mul_func = root.mul.get_concrete_function(input_data)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(
|
|
root, save_dir, {'add': add_func, 'sub': sub_func, 'mul': mul_func}
|
|
)
|
|
|
|
# Try converting multiple functions.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Make sure that the weight tensors are shared.
|
|
self.assertLess(len(tflite_model), 1100000)
|
|
|
|
interp = lite.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
self.assertLen(signature_defs, 3)
|
|
add_signature_runner = interp.get_signature_runner('add')
|
|
sub_signature_runner = interp.get_signature_runner('sub')
|
|
mul_signature_runner = interp.get_signature_runner('mul')
|
|
self.assertIsNotNone(add_signature_runner)
|
|
self.assertIsNotNone(sub_signature_runner)
|
|
self.assertIsNotNone(mul_signature_runner)
|
|
|
|
@test_util.run_v2_only
|
|
def testNoConcreteFunctionModel(self):
|
|
root = self._getMultiFunctionModel()
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir)
|
|
|
|
with self.assertRaises(ValueError) as error:
|
|
_ = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
self.assertIn(
|
|
'Only support at least one signature key.', str(error.exception)
|
|
)
|
|
|
|
@test_util.run_v2_only
|
|
def testKerasSequentialModel(self):
|
|
"""Test a simple sequential tf.Keras model."""
|
|
input_data = tf.constant(1.0, shape=[1, 1])
|
|
|
|
x = np.array([[1.0], [2.0]])
|
|
y = np.array([[2.0], [4.0]])
|
|
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.layers.Dropout(0.2),
|
|
tf.keras.layers.Dense(1),
|
|
])
|
|
model.compile(optimizer='sgd', loss='mean_squared_error')
|
|
model.fit(x, y, epochs=1)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(model, save_dir)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = model.predict(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testKerasSequentialModelExport(self):
|
|
"""Test a simple sequential tf.Keras model with `model.export` usage."""
|
|
input_data = tf.constant(1.0, shape=[1, 1])
|
|
|
|
x = np.array([[1.0], [2.0]])
|
|
y = np.array([[2.0], [4.0]])
|
|
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.layers.Dropout(0.2),
|
|
tf.keras.layers.Dense(1),
|
|
])
|
|
model.compile(optimizer='sgd', loss='mean_squared_error')
|
|
model.fit(x, y, epochs=1)
|
|
|
|
export_dir = os.path.join(self.get_temp_dir(), 'exported_model')
|
|
model.export(export_dir)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(export_dir)
|
|
tflite_model = converter.convert()
|
|
|
|
# Validate endpoints following `.export` to TFLite conversion.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
self.assertLen(signature_defs, 1)
|
|
self.assertEqual(next(iter(signature_defs)), 'serving_default')
|
|
|
|
# Check values from converted model.
|
|
expected_value = model.predict(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testGraphDebugInfo(self):
|
|
"""Test a SavedModel has debug info captured."""
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = tf.function(lambda x: 2.0 * x)
|
|
to_save = root.f.get_concrete_function(input_data)
|
|
options = save_options.SaveOptions(save_debug_info=True)
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, to_save, options)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
converter.convert()
|
|
self._assertValidDebugInfo(converter._debug_info)
|
|
|
|
@test_util.run_v2_only
|
|
def testNonStatefulConvLSTM2D(self):
|
|
"""Test saved model with non stateful ConvLSTM2D keras layer."""
|
|
# Create keras model
|
|
model = tf.keras.Sequential([
|
|
tf.keras.layers.ConvLSTM2D(
|
|
32,
|
|
(3, 3),
|
|
padding='same',
|
|
return_sequences=True,
|
|
stateful=False,
|
|
batch_input_shape=(1, 1, 10, 10, 1),
|
|
)
|
|
])
|
|
model.compile()
|
|
|
|
# Export the keras model to saved model.
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'conv_lstm_2d')
|
|
model.save(saved_model_dir, save_format='tf', include_optimizer=False)
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
@test_util.run_v2_only
|
|
def testKerasConvLSTM2DWithMoreThanOneDilationRate(self):
|
|
input_tensor = tf.keras.layers.Input(
|
|
batch_size=8,
|
|
shape=[9, 10, 11, 12],
|
|
name='input_tensor',
|
|
dtype=tf.float32,
|
|
)
|
|
|
|
output = tf.keras.layers.ConvLSTM2D(
|
|
filters=3,
|
|
kernel_size=3,
|
|
strides=1,
|
|
padding='VALID',
|
|
dilation_rate=2,
|
|
use_bias=False,
|
|
bias_initializer='ones',
|
|
data_format='channels_last',
|
|
)(input_tensor)
|
|
|
|
model = tf.keras.Model(inputs=[input_tensor], outputs=output)
|
|
model.compile(
|
|
optimizer='adam',
|
|
loss='sparse_categorical_crossentropy',
|
|
metrics=['accuracy'],
|
|
)
|
|
|
|
# Export the keras model to saved model.
|
|
saved_model_dir = os.path.join(
|
|
self.get_temp_dir(), 'conv_lstm_2d_with_dilation_rate'
|
|
)
|
|
model.save(saved_model_dir, save_format='tf', include_optimizer=False)
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
@test_util.run_v2_only
|
|
def testKerasFullyConnectedOutputShape3D(self):
|
|
"""Create a simple FullyConnected Model with an output of three dimensions."""
|
|
input_tensor = tf.keras.layers.Input(
|
|
batch_size=1, shape=[3, 3], name='input_tensor', dtype=tf.float32
|
|
)
|
|
|
|
class _FakeQuantArgsLayer(tf.keras.layers.Layer):
|
|
"""A fake quantization layer with fake_quant_with_min_max_args.
|
|
|
|
Keras 3 requires wrapping the tf function inside Keras layer.
|
|
"""
|
|
|
|
def call(self, x):
|
|
return tf.quantization.fake_quant_with_min_max_args(x, -3.0, 3.0)
|
|
|
|
x = _FakeQuantArgsLayer()(input_tensor)
|
|
x = tf.keras.layers.Dense(3)(x)
|
|
x = _FakeQuantArgsLayer()(x)
|
|
model = tf.keras.Model(input_tensor, x)
|
|
|
|
model.compile(
|
|
optimizer='adam', loss='mean_squared_error', metrics=['accuracy']
|
|
)
|
|
|
|
# Export the keras model to saved model.
|
|
saved_model_dir = os.path.join(
|
|
self.get_temp_dir(), 'fully_connected_output_3d'
|
|
)
|
|
model.save(saved_model_dir, save_format='tf', include_optimizer=False)
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
output_details = interp.get_output_details()
|
|
input_details = interp.get_input_details()
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]], np.float32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
interp.invoke()
|
|
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
expected_value = model.predict(input_data)
|
|
|
|
self.assertLen(output_details[0]['shape_signature'], 3)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-1)
|
|
self.assertEqual(
|
|
list(output_details[0]['shape_signature']),
|
|
list(model.layers[-1].output_shape),
|
|
)
|
|
|
|
@test_util.run_v2_only
|
|
def testKerasConv2DTransposedWithMismatchQuantizedAxes(self):
|
|
class QuantConv2DTransposed(tf.keras.layers.Layer):
|
|
|
|
def build(self, input_shape):
|
|
self.kernel = self.add_weight('kernel', [3, 3, input_shape[-1], 24])
|
|
|
|
def call(self, inputs):
|
|
filters = tf.quantization.fake_quant_with_min_max_vars_per_channel(
|
|
self.kernel,
|
|
-3.0 * tf.ones([24]),
|
|
3.0 * tf.ones([24]),
|
|
narrow_range=True,
|
|
)
|
|
filters = tf.transpose(filters, (0, 1, 3, 2))
|
|
return tf.nn.conv2d_transpose(
|
|
inputs, filters, [*inputs.shape[:-1], 24], 1
|
|
)
|
|
|
|
class _FakeQuantVarsLayer(tf.keras.layers.Layer):
|
|
"""A fake quantization layer with fake_quant_with_min_max_vars.
|
|
|
|
Keras 3 requires wrapping the tf function inside Keras layer.
|
|
"""
|
|
|
|
def call(self, x):
|
|
return tf.quantization.fake_quant_with_min_max_vars(
|
|
x, -3.0, 3.0, narrow_range=True)
|
|
|
|
inp = tf.keras.Input(shape=(6, 8, 48), batch_size=1)
|
|
x = _FakeQuantVarsLayer()(inp)
|
|
x = QuantConv2DTransposed()(x)
|
|
x = _FakeQuantVarsLayer()(x)
|
|
|
|
model = tf.keras.Model(inp, x)
|
|
|
|
saved_model_dir = os.path.join(
|
|
self.get_temp_dir(), 'keras_conv2d_transpose'
|
|
)
|
|
model.save(saved_model_dir)
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
|
|
with self.assertRaises(convert.ConverterError) as error:
|
|
_ = converter.convert()
|
|
self.assertIn(
|
|
'mismatched quantized axes of input and output', str(error.exception)
|
|
)
|
|
|
|
def _createModelWithInputShape(self, shape):
|
|
"""Create a simple SavedModel with a certain shape."""
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'input_shape_model')
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
unknown_shape = tf.TensorShape(shape)
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=unknown_shape, dtype=tf.float32, name='input'
|
|
)
|
|
out_tensor = in_tensor + in_tensor
|
|
inputs = {'input': in_tensor}
|
|
outputs = {'output': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
@test_util.run_v2_only
|
|
def testUnknownInputShapeModel(self):
|
|
"""Test a SavedModel with an unknown input shape."""
|
|
saved_model_dir = self._createModelWithInputShape(None)
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
# Validate that tensors with unknown shape have unknown rank.
|
|
tflite_model_obj = flatbuffer_utils.convert_bytearray_to_object(
|
|
tflite_model
|
|
)
|
|
for tensor in tflite_model_obj.subgraphs[0].tensors:
|
|
self.assertEqual(False, tensor.hasRank)
|
|
self.assertEqual([], tensor.shape.tolist())
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
input_data = np.array([1.0, 2.0, 3.0], dtype=np.float32)
|
|
interp.resize_tensor_input(input_details[0]['index'], [3], strict=False)
|
|
interp.allocate_tensors()
|
|
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual([2.0, 4.0, 6.0], list(actual_value))
|
|
|
|
@test_util.run_v2_only
|
|
def testScalarInputShapeModel(self):
|
|
"""Test a SavedModel with a scalar input."""
|
|
saved_model_dir = self._createModelWithInputShape([])
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
# Validate that scalar tensors have a rank = 0.
|
|
tflite_model_obj = flatbuffer_utils.convert_bytearray_to_object(
|
|
tflite_model
|
|
)
|
|
for tensor in tflite_model_obj.subgraphs[0].tensors:
|
|
self.assertEqual(True, tensor.hasRank)
|
|
self.assertEqual([], tensor.shape.tolist())
|
|
|
|
@test_util.run_v2_only
|
|
def testMatrixInputShapeModel(self):
|
|
"""Test a SavedModel with a matrix input."""
|
|
saved_model_dir = self._createModelWithInputShape([2, 3])
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
# Validate that matrix tensors have a rank = 2.
|
|
tflite_model_obj = flatbuffer_utils.convert_bytearray_to_object(
|
|
tflite_model
|
|
)
|
|
for tensor in tflite_model_obj.subgraphs[0].tensors:
|
|
self.assertEqual(True, tensor.hasRank)
|
|
self.assertEqual([2, 3], tensor.shape.tolist())
|
|
|
|
@parameterized.named_parameters(
|
|
('_PerChannelQuant', False, False),
|
|
('_PerChannelMlirQuant', False, True),
|
|
('_PerTensorQuant', True, False),
|
|
('_PerTensorMlirQuant', True, True),
|
|
('_PerChannelDynamicRange', False, False, True),
|
|
('_PerTensorDynamicRange', True, False, True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testDisablePerChannelQuantization(
|
|
self,
|
|
disable_per_channel=False,
|
|
enable_mlir_quantizer=False,
|
|
representative_dataset=True,
|
|
):
|
|
# Dynamic range quant requires total num elements of filters > 1024.
|
|
k_num_filters = 38
|
|
model = tf.keras.models.Sequential(
|
|
[tf.keras.layers.Conv2D(k_num_filters, (3, 3), activation='relu')]
|
|
)
|
|
model.build(input_shape=(1, 5, 5, 3))
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'conv_saved_model')
|
|
save.save(model, saved_model_dir)
|
|
k_conv_name = 'sequential/conv2d/Conv2D'
|
|
quantized_converter = lite.TFLiteConverterV2.from_saved_model(
|
|
saved_model_dir
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
if representative_dataset:
|
|
|
|
def calib_gen():
|
|
for _ in range(5):
|
|
yield [np.random.uniform(-1, 1, size=(1, 5, 5, 3)).astype(np.float32)]
|
|
|
|
quantized_converter.representative_dataset = calib_gen
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS
|
|
]
|
|
quantized_converter.experimental_new_quantizer = enable_mlir_quantizer
|
|
if disable_per_channel:
|
|
quantized_converter._experimental_disable_per_channel = (
|
|
disable_per_channel
|
|
)
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
detail = next((
|
|
d
|
|
for d in interp.get_tensor_details()
|
|
if d['name'].startswith(k_conv_name)
|
|
))
|
|
quant_params = detail['quantization_parameters']
|
|
expected_num_params = k_num_filters
|
|
if disable_per_channel:
|
|
expected_num_params = 1
|
|
self.assertLen(quant_params['scales'], expected_num_params)
|
|
if len(quant_params['zero_points']) != 1:
|
|
self.assertLen(quant_params['zero_points'], expected_num_params)
|
|
|
|
@parameterized.named_parameters(
|
|
('_INT8Quant_INT32Bias', False, False, dtypes.int32, True),
|
|
('_INT16Quant_INT64Bias', True, False, dtypes.int64, True),
|
|
('_INT8Quant_INT32Bias_Set', False, True, dtypes.int32, True),
|
|
('_INT8Quant_INT64Bias_Set', False, True, dtypes.int64, False),
|
|
('_INT16Quant_INT32Bias_Set', True, True, dtypes.int32, True),
|
|
('_INT16Quant_INT64Bias_Set', True, True, dtypes.int64, True),
|
|
('_INT16Quant_FLOAT32Bias_Set', True, True, dtypes.float32, False),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testBiasQuantization(
|
|
self,
|
|
is_int16_quantize,
|
|
explicitly_set_bias,
|
|
bias_type,
|
|
is_valid_bias_type,
|
|
):
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.layers.Dense(
|
|
1024, input_shape=[1024], activation=None, bias_initializer='ones'
|
|
)
|
|
])
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'dense_saved_model')
|
|
save.save(model, saved_model_dir)
|
|
k_dense_bias_name = 'sequential/dense/BiasAdd'
|
|
quantized_converter = lite.TFLiteConverterV2.from_saved_model(
|
|
saved_model_dir
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
|
|
if explicitly_set_bias:
|
|
quantized_converter._experimental_full_integer_quantization_bias_type = (
|
|
bias_type
|
|
)
|
|
|
|
if is_int16_quantize:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
|
|
]
|
|
else:
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8
|
|
]
|
|
|
|
def calibration_gen():
|
|
for _ in range(5):
|
|
yield [np.random.randn(1, 1024).astype(np.float32)]
|
|
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
|
|
if not is_valid_bias_type:
|
|
with self.assertRaisesRegex(ValueError, 'Expected bias type to be'):
|
|
quantized_converter.convert()
|
|
return
|
|
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
dense_bias = next((
|
|
d
|
|
for d in interp.get_tensor_details()
|
|
if d['name'].startswith(k_dense_bias_name)
|
|
))
|
|
self.assertEqual(bias_type, dense_bias['dtype'])
|
|
|
|
@parameterized.named_parameters(
|
|
('_Int8PerChannelMlirDynamicRangeQuant', True, False, False),
|
|
('_Int8PerChannelTocoDynamicRangeQuant', False, False, False),
|
|
('_Int8PerTensorMlirDynamicRangeQuant', True, True, False),
|
|
('_Int8PerTensorTocoDynamicRangeQuant', False, True, False),
|
|
('_Float16DynamicRangeQuant', True, False, True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testMlirDynamicRangeQuantization(
|
|
self,
|
|
enable_new_dynamic_range_quantizer,
|
|
disable_per_channel,
|
|
enable_float16_quant,
|
|
):
|
|
num_filters = 1024
|
|
conv_name = 'sequential/conv2d/Conv2D'
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.layers.Conv2D(
|
|
num_filters, (3, 3), activation='relu', bias_initializer='ones'
|
|
)
|
|
])
|
|
model.build(input_shape=(1, 32, 32, 3))
|
|
saved_model_dir = self.create_tempdir()
|
|
save.save(model, saved_model_dir.full_path)
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(
|
|
saved_model_dir.full_path
|
|
)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.experimental_new_dynamic_range_quantizer = (
|
|
enable_new_dynamic_range_quantizer
|
|
)
|
|
converter._experimental_disable_per_channel = disable_per_channel
|
|
if enable_float16_quant:
|
|
converter.target_spec.supported_types = [tf.float16]
|
|
quantized_tflite_model = converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
# Do not apply delegates as XNNPack converts per tensor to per channel.
|
|
interp = interpreter.Interpreter(
|
|
model_content=quantized_tflite_model,
|
|
experimental_op_resolver_type=interpreter.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES,
|
|
)
|
|
interp.allocate_tensors()
|
|
quantized_weight = None
|
|
quantized_weight_with_one_postfix = None
|
|
quantized_weight_without_one_postfix = None
|
|
for d in interp.get_tensor_details():
|
|
if d['name'] == conv_name + '1':
|
|
quantized_weight = d
|
|
quantized_weight_with_one_postfix = d
|
|
break
|
|
for d in interp.get_tensor_details():
|
|
if d['name'].startswith(conv_name):
|
|
if quantized_weight is None:
|
|
quantized_weight = d
|
|
quantized_weight_without_one_postfix = d
|
|
break
|
|
|
|
self.assertIsNotNone(quantized_weight)
|
|
quant_params = quantized_weight['quantization_parameters']
|
|
|
|
if enable_float16_quant:
|
|
expected_num_params = 0
|
|
else:
|
|
expected_num_params = 1 if disable_per_channel else num_filters
|
|
self.assertLen(quant_params['scales'], expected_num_params)
|
|
if len(quant_params['zero_points']) != 1:
|
|
self.assertLen(quant_params['zero_points'], expected_num_params)
|
|
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
self.assertEqual(np.float32, input_details[0]['dtype'])
|
|
self.assertEqual(np.float32, output_details[0]['dtype'])
|
|
if enable_float16_quant:
|
|
self.assertTrue(
|
|
(
|
|
quantized_weight_with_one_postfix is not None
|
|
and np.float16 == quantized_weight_with_one_postfix['dtype']
|
|
)
|
|
or (
|
|
quantized_weight_without_one_postfix is not None
|
|
and np.float16 == quantized_weight_without_one_postfix['dtype']
|
|
)
|
|
)
|
|
else:
|
|
self.assertEqual(np.int8, quantized_weight['dtype'])
|
|
|
|
@parameterized.named_parameters(
|
|
('_NONE', 'NONE'),
|
|
('_STATIC', 'STATIC'),
|
|
('_DYNAMIC', 'DYNAMIC'),
|
|
('_UNKNOWN', 'UNKNOWN'),
|
|
)
|
|
def testQDQConversionMode(self, mode):
|
|
num_filters = 1024
|
|
model = tf.keras.models.Sequential(
|
|
[tf.keras.layers.Conv2D(num_filters, (3, 3), activation='relu')]
|
|
)
|
|
model.build(input_shape=(1, 32, 32, 3))
|
|
saved_model_dir = self.create_tempdir()
|
|
save.save(model, saved_model_dir.full_path)
|
|
converter = lite.TFLiteConverterV2.from_saved_model(
|
|
saved_model_dir.full_path
|
|
)
|
|
converter._experimental_qdq_conversion_mode = mode
|
|
|
|
if mode == 'UNKNOWN':
|
|
with self.assertRaises(convert.ConverterError) as error:
|
|
converter.convert()
|
|
self.assertIn('Unknown QDQ conversion mode:', str(error.exception))
|
|
else:
|
|
model = converter.convert()
|
|
self.assertIsNotNone(model)
|
|
|
|
|
|
class FromKerasModelTest(lite_v2_test_util.ModelTest):
|
|
|
|
@parameterized.named_parameters(
|
|
('EnableMlirVariableQuantizationNumState1', True, 1),
|
|
('DisablMlirVariableQuantizationNumState1', False, 1),
|
|
('EnableMlirVariableQuantizationNumState2', True, 2),
|
|
('DisablMlirVariableQuantizationNumState2', False, 2),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testVariableQuantization(self, variable_quantization, number_of_states):
|
|
k_readvariable_name = 'model/read_assign/concat/ReadVariableOp'
|
|
model, calibration_gen = self._createReadAssignModel(number_of_states)
|
|
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = calibration_gen
|
|
converter.target_spec.supported_ops = [lite.OpsSet.TFLITE_BUILTINS_INT8]
|
|
converter.inference_input_type = tf.int8 # or tf.uint8
|
|
converter.inference_output_type = tf.int8 # or tf.uint8
|
|
converter._experimental_variable_quantization = variable_quantization
|
|
|
|
quantized_tflite_model = converter.convert()
|
|
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
|
|
detail = next((
|
|
d
|
|
for d in interp.get_tensor_details()
|
|
if d['name'].startswith(k_readvariable_name)
|
|
))
|
|
quant_params = detail['quantization_parameters']
|
|
if variable_quantization:
|
|
expected_num_params = 1
|
|
else:
|
|
# This number is not a spec. Since It's the unintended number, it can be
|
|
# changed later by the other features of the quantizer.
|
|
expected_num_params = 0
|
|
|
|
self.assertLen(quant_params['scales'], expected_num_params)
|
|
self.assertLen(quant_params['zero_points'], expected_num_params)
|
|
|
|
@parameterized.named_parameters(
|
|
('EnableMlirVariableQuantizationNumState1', True, 1),
|
|
('DisablMlirVariableQuantizationNumState1', False, 1),
|
|
('EnableMlirVariableQuantizationNumState2', True, 2),
|
|
('DisablMlirVariableQuantizationNumState2', False, 2),
|
|
)
|
|
def testVariableQuantizationInFloat16(
|
|
self, variable_quantization, number_of_states
|
|
):
|
|
model, _ = self._createReadAssignModel(number_of_states)
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.target_spec.supported_types = [tf.float16]
|
|
converter._experimental_variable_quantization = variable_quantization
|
|
|
|
if variable_quantization:
|
|
with self.assertRaises(ValueError) as error:
|
|
converter.convert()
|
|
self.assertIn(
|
|
'`_experimental_variable_quantization` is only supported for full',
|
|
str(error.exception),
|
|
)
|
|
else:
|
|
converter.convert()
|
|
|
|
@test_util.run_v2_only
|
|
def testSequentialModel(self):
|
|
"""Test a simple sequential tf.Keras model."""
|
|
input_data = tf.constant(1.0, shape=[1, 1])
|
|
|
|
# Create a simple Keras model.
|
|
x = np.array([[1.0], [2.0]])
|
|
y = np.array([[2.0], [4.0]])
|
|
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.layers.Dropout(0.2),
|
|
tf.keras.layers.Dense(units=1, input_shape=[1]),
|
|
])
|
|
model.compile(optimizer='sgd', loss='mean_squared_error')
|
|
model.fit(x, y, epochs=1)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
tflite_model = converter.convert()
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(
|
|
metadata.environment.modelType, metadata_fb.ModelType.KERAS_MODEL
|
|
)
|
|
|
|
# Check values from converted model.
|
|
expected_value = model.predict(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testSequentialMultiInputOutputModel(self):
|
|
"""Test a tf.Keras model with multiple inputs and outputs."""
|
|
left_input_data = tf.constant(1.0, shape=[1, 3])
|
|
right_input_data = tf.constant(1.0, shape=[1, 3])
|
|
|
|
# Create a simple Keras model.
|
|
input_a_np = np.random.random((10, 3))
|
|
input_b_np = np.random.random((10, 3))
|
|
output_c_np = np.random.random((10, 3))
|
|
output_d_np = np.random.random((10, 2))
|
|
|
|
input_a = tf.keras.layers.Input(shape=(3,), name='input_a')
|
|
input_b = tf.keras.layers.Input(shape=(3,), name='input_b')
|
|
|
|
dense = tf.keras.layers.Dense(8, name='dense_1')
|
|
interm_a = dense(input_a)
|
|
interm_b = dense(input_b)
|
|
merged = tf.keras.layers.concatenate([interm_a, interm_b], name='merge')
|
|
|
|
output_c = tf.keras.layers.Dense(3, activation='softmax', name='dense_2')(
|
|
merged
|
|
)
|
|
output_d = tf.keras.layers.Dense(2, activation='softmax', name='dense_3')(
|
|
merged
|
|
)
|
|
|
|
model = tf.keras.models.Model(
|
|
inputs=[input_a, input_b], outputs=[output_c, output_d]
|
|
)
|
|
model.compile(optimizer='sgd', loss='mean_squared_error')
|
|
model.fit([input_a_np, input_b_np], [output_c_np, output_d_np], epochs=1)
|
|
|
|
# Convert model and ensure model is not None.
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
input_data = [left_input_data, right_input_data]
|
|
expected_value = model.predict(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, input_data)
|
|
for tf_result, tflite_result in zip(expected_value, actual_value):
|
|
self.assertAllClose(tf_result, tflite_result, atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testGraphDebugInfo(self):
|
|
"""Test a tf.Keras model has debug info captured."""
|
|
# Create a simple Keras model.
|
|
x = [-1, 0, 1, 2, 3, 4]
|
|
y = [-3, -1, 1, 3, 5, 7]
|
|
model = tf.keras.models.Sequential(
|
|
[tf.keras.layers.Dense(units=1, input_shape=[1])]
|
|
)
|
|
model.compile(optimizer='sgd', loss='mean_squared_error')
|
|
model.fit(x, y, epochs=1)
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter.convert()
|
|
self._assertValidDebugInfo(converter._debug_info)
|
|
|
|
@test_util.run_v2_only
|
|
def testKerasFallbackPath(self):
|
|
"""Test keras model which failed when exporting to the saved model."""
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((20)), dtype=np.float32)
|
|
)
|
|
|
|
class Model(tf.keras.Model):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
# A None name will cause a failure in exporting to a saved model.
|
|
self.shared_weights = self.add_weight(
|
|
name=None,
|
|
shape=(20, 1),
|
|
dtype=tf.float32,
|
|
initializer=tf.random_normal_initializer(
|
|
mean=0.0, stddev=300 ** (-0.5)
|
|
),
|
|
)
|
|
|
|
def call(self, x):
|
|
return tf.add(self.shared_weights, x)
|
|
|
|
# Building the model.
|
|
model = Model()
|
|
model.compile(optimizer='sgd', loss='mean_squared_error')
|
|
model.fit(input_data, input_data, epochs=1)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
|
|
@test_util.run_v2_only
|
|
def testSignatureDefs(self):
|
|
"""Test converting SignatureDef is correct and uses SignatureDef API."""
|
|
keras_model = tf.keras.Sequential([
|
|
tf.keras.layers.Conv2D(
|
|
32,
|
|
kernel_size=3,
|
|
padding='same',
|
|
activation='relu',
|
|
input_shape=(32, 32, 3),
|
|
name='tensor',
|
|
),
|
|
tf.keras.layers.Dense(10, name='output_tensor'),
|
|
])
|
|
|
|
converter = lite.TFLiteConverterV2.from_keras_model(keras_model)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
input_data = tf.constant(
|
|
np.random.uniform(-1, 1, size=(1, 32, 32, 3)).astype(np.float32)
|
|
)
|
|
expected_value = keras_model(input_data)
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
results = self._evaluateTFLiteModelUsingSignatureDef(
|
|
tflite_model, 'serving_default', {'tensor_input': input_data}
|
|
)
|
|
self.assertEqual(list(results.keys()), ['output_tensor'])
|
|
self.assertAllClose(expected_value.numpy(), results['output_tensor'])
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 1)
|
|
self.assertEqual(list(signature_defs.keys()), ['serving_default'])
|
|
self.assertLen(signature_defs.values(), 1)
|
|
self.assertEqual(
|
|
list(signature_defs['serving_default'].keys()), ['inputs', 'outputs']
|
|
)
|
|
self.assertCountEqual(
|
|
signature_defs['serving_default']['inputs'], ['tensor_input']
|
|
)
|
|
self.assertEqual(
|
|
list(signature_defs['serving_default']['outputs']), ['output_tensor']
|
|
)
|
|
|
|
@parameterized.named_parameters(
|
|
('_PerChannelMlirDynamicRangeQuant', True, False, False),
|
|
('_PerChannelTocoDynamicRangeQuant', False, False, False),
|
|
('_PerTensorMlirDynamicRangeQuant', True, True, False),
|
|
('_PerTensorTocoDynamicRangeQuant', False, True, False),
|
|
('_Float16DynamicRangeQuant', True, False, True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testMlirDynamicRangeQuantization(
|
|
self,
|
|
enable_new_dynamic_range_quantizer,
|
|
disable_per_channel,
|
|
enable_float16_quant,
|
|
):
|
|
num_filters = 1024
|
|
conv_name = 'sequential/conv2d/Conv2D'
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.Input(shape=(32, 32, 3)),
|
|
tf.keras.layers.Conv2D(
|
|
num_filters, (3, 3), activation='relu', bias_initializer='ones'
|
|
),
|
|
])
|
|
model.build()
|
|
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.experimental_new_dynamic_range_quantizer = (
|
|
enable_new_dynamic_range_quantizer
|
|
)
|
|
converter._experimental_disable_per_channel = disable_per_channel
|
|
if enable_float16_quant:
|
|
converter.target_spec.supported_types = [tf.float16]
|
|
quantized_tflite_model = converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
# Do not apply delegates as XNNPack converts per tensor to per channel.
|
|
interp = interpreter.Interpreter(
|
|
model_content=quantized_tflite_model,
|
|
experimental_op_resolver_type=interpreter.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES,
|
|
)
|
|
interp.allocate_tensors()
|
|
quantized_weight = None
|
|
quantized_weight_with_one_postfix = None
|
|
quantized_weight_without_one_postfix = None
|
|
for d in interp.get_tensor_details():
|
|
if d['name'] == conv_name + '1':
|
|
quantized_weight = d
|
|
quantized_weight_with_one_postfix = d
|
|
break
|
|
for d in interp.get_tensor_details():
|
|
if d['name'].startswith(conv_name):
|
|
if quantized_weight is None:
|
|
quantized_weight = d
|
|
quantized_weight_without_one_postfix = d
|
|
break
|
|
|
|
self.assertIsNotNone(quantized_weight)
|
|
quant_params = quantized_weight['quantization_parameters']
|
|
|
|
if enable_float16_quant:
|
|
expected_num_params = 0
|
|
else:
|
|
expected_num_params = 1 if disable_per_channel else num_filters
|
|
self.assertLen(quant_params['scales'], expected_num_params)
|
|
if len(quant_params['zero_points']) != 1:
|
|
self.assertLen(quant_params['zero_points'], expected_num_params)
|
|
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
self.assertEqual(np.float32, input_details[0]['dtype'])
|
|
self.assertEqual(np.float32, output_details[0]['dtype'])
|
|
if enable_float16_quant:
|
|
self.assertTrue(
|
|
(
|
|
quantized_weight_with_one_postfix is not None
|
|
and np.float16 == quantized_weight_with_one_postfix['dtype']
|
|
)
|
|
or (
|
|
quantized_weight_without_one_postfix is not None
|
|
and np.float16 == quantized_weight_without_one_postfix['dtype']
|
|
)
|
|
)
|
|
else:
|
|
self.assertEqual(np.int8, quantized_weight['dtype'])
|
|
|
|
# TODO(b/242081598): The num_bits parameter should be restored to (2, 4, 6)
|
|
# once a 4-bit conv kernel is available.
|
|
@parameterized.named_parameters([
|
|
(
|
|
'{}BitWeightOnly={}LowBit={}'.format(num_bits, weight_only, low_bit),
|
|
num_bits,
|
|
weight_only,
|
|
low_bit,
|
|
)
|
|
for num_bits, weight_only, low_bit in itertools.product(
|
|
(5, 7, 6), (True, False), (True, False)
|
|
)
|
|
])
|
|
@test_util.run_v2_only
|
|
def testQATLowBitKerasModel(self, num_bits, weight_only, low_bit):
|
|
bit_max = (1 << (num_bits - 1)) - 1
|
|
bit_min = -bit_max
|
|
tf_input_shape = (5, 5, 3)
|
|
tflite_input_shape = (1,) + tf_input_shape
|
|
model, input_name, output_name = self._createV2QATLowBitKerasModel(
|
|
tf_input_shape, weight_only, num_bits, bit_min, bit_max
|
|
)
|
|
input_data = np.linspace(0, 6, np.prod(tflite_input_shape)).reshape(
|
|
tflite_input_shape
|
|
)
|
|
tf_result = model(input_data)
|
|
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
if low_bit:
|
|
converter._experimental_low_bit_qat = True
|
|
tflite_model = converter.convert()
|
|
result = self._evaluateTFLiteModelUsingSignatureDef(
|
|
tflite_model,
|
|
'serving_default',
|
|
{input_name: input_data.astype(np.float32)},
|
|
)[output_name]
|
|
self.assertAllClose(
|
|
[np.linalg.norm(result - tf_result.numpy().astype(np.float32))], [0.0]
|
|
)
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
interp.allocate_tensors()
|
|
num_8bit_activations = 0
|
|
num_8bit_weights = 0
|
|
kernel_name = (
|
|
'model/conv_wrapper/Conv2D;model/conv_wrapper/'
|
|
'FakeQuantWithMinMaxVarsPerChannel'
|
|
)
|
|
|
|
for detail in interp.get_tensor_details():
|
|
if (
|
|
detail['dtype'] == np.int8
|
|
and detail['name']
|
|
and detail['name'] == kernel_name
|
|
):
|
|
num_8bit_weights += 1
|
|
weights = interp.get_tensor(detail['index'])
|
|
if low_bit:
|
|
self.assertFalse(
|
|
(bit_min > weights).any() or (weights > bit_max).any()
|
|
)
|
|
else:
|
|
self.assertTrue(
|
|
(bit_min > weights).any() or (weights > bit_max).any()
|
|
)
|
|
self.assertIn('scales', detail['quantization_parameters'])
|
|
if low_bit and detail['quantization_parameters']['scales']:
|
|
self.assertAllClose(
|
|
detail['quantization_parameters']['scales'], [1.0]
|
|
)
|
|
elif detail['dtype'] == np.int8 and detail['name']:
|
|
self.assertFalse(weight_only)
|
|
self.assertIn('scales', detail['quantization_parameters'])
|
|
if detail['quantization_parameters']['scales']:
|
|
self.assertAllClose(
|
|
detail['quantization_parameters']['scales'], [6 / 255]
|
|
)
|
|
num_8bit_activations += 1
|
|
|
|
self.assertEqual(num_8bit_weights, 0 if weight_only and not low_bit else 1)
|
|
# 3 activations with full integer: conv_input, conv_output, reshape_output
|
|
self.assertEqual(num_8bit_activations, 0 if weight_only else 3)
|
|
|
|
@test_util.run_v2_only
|
|
def testKerasConv2DTransposedWithBiasAndActivation(self):
|
|
class QuantConv2DTransposedWithBiasAndActivation(tf.keras.layers.Layer):
|
|
|
|
def build(self, input_shape):
|
|
self.kernel = self.add_weight('kernel', (3, 3, input_shape[-1], 3))
|
|
self.bias = self.add_weight('bias', (3,))
|
|
|
|
def call(self, inputs):
|
|
filters = tf.quantization.fake_quant_with_min_max_vars(
|
|
self.kernel, -3.0, 3.0, narrow_range=True
|
|
)
|
|
filters = tf.transpose(filters, (0, 1, 3, 2))
|
|
result = tf.nn.conv2d_transpose(
|
|
inputs, filters, [*inputs.shape[:-1], 3], 1
|
|
)
|
|
result = tf.nn.bias_add(result, self.bias)
|
|
result = tf.nn.relu(result)
|
|
|
|
return tf.quantization.fake_quant_with_min_max_vars(
|
|
result, -3.0, 3.0, narrow_range=True
|
|
)
|
|
|
|
class _FakeQuantVarsLayer(tf.keras.layers.Layer):
|
|
"""A fake quantization layer with fake_quant_with_min_max_vars.
|
|
|
|
Keras 3 requires wrapping the tf function inside Keras layer.
|
|
"""
|
|
|
|
def call(self, x):
|
|
return tf.quantization.fake_quant_with_min_max_vars(
|
|
x, -3.0, 3.0, narrow_range=True)
|
|
|
|
inp = tf.keras.Input(shape=(6, 8, 6), batch_size=1)
|
|
x = _FakeQuantVarsLayer()(inp)
|
|
x = QuantConv2DTransposedWithBiasAndActivation()(x)
|
|
|
|
model = tf.keras.Model(inp, x)
|
|
|
|
tf_input_shape = (1, 6, 8, 6)
|
|
input_data = np.linspace(0, 6, np.prod(tf_input_shape)).reshape(
|
|
tf_input_shape
|
|
)
|
|
tf_result = model(input_data)
|
|
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
|
|
converted_model = converter.convert()
|
|
tf.lite.experimental.Analyzer.analyze(model_content=converted_model)
|
|
|
|
interp = interpreter.Interpreter(model_content=converted_model)
|
|
interp.allocate_tensors()
|
|
|
|
input_index = interp.get_input_details()[0]['index']
|
|
output_index = interp.get_output_details()[0]['index']
|
|
|
|
interp.set_tensor(input_index, input_data.astype(np.float32))
|
|
interp.invoke()
|
|
tflite_result = interp.tensor(output_index)()
|
|
|
|
self.assertAllClose(
|
|
[np.linalg.norm(tflite_result - tf_result.numpy().astype(np.float32))],
|
|
[0.0],
|
|
)
|
|
|
|
num_float32_tensor = 0
|
|
for detail in interp.get_tensor_details():
|
|
if detail['dtype'] == np.float32:
|
|
num_float32_tensor += 1
|
|
|
|
# There should be only 2 float tensors, input and output.
|
|
self.assertEqual(num_float32_tensor, 2)
|
|
|
|
@parameterized.named_parameters(
|
|
('_PerChannelQuant', False, False),
|
|
('_PerChannelMlirQuant', False, True),
|
|
('_PerTensorQuant', True, False),
|
|
('_PerTensorMlirQuant', True, True),
|
|
('_PerChannelDynamicRange', False, True, True),
|
|
('_PerTensorDynamicRange', True, True, True),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testDisablePerChannelQuantizationForDenseLayers(
|
|
self,
|
|
disable_per_channel_for_dense=False,
|
|
enable_mlir_quantizer=False,
|
|
representative_dataset=False,
|
|
):
|
|
k_dense_name = 'MatMul'
|
|
# Dynamic range quant requires total num elements of filters > 1024.
|
|
k_num_filters = 64
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.Input(shape=(16,)),
|
|
tf.keras.layers.Dense(k_num_filters, activation='relu'),
|
|
])
|
|
model.build()
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
if representative_dataset:
|
|
|
|
def calibration_gen():
|
|
for _ in range(5):
|
|
yield [np.random.uniform(-1, 1, size=(1, 16)).astype(np.float32)]
|
|
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
quantized_converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS
|
|
]
|
|
quantized_converter.experimental_new_quantizer = enable_mlir_quantizer
|
|
if disable_per_channel_for_dense:
|
|
quantized_converter._experimental_disable_per_channel_quantization_for_dense_layers = (
|
|
disable_per_channel_for_dense
|
|
)
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
# Do not apply delegates as XNNPack converts per tensor to per channel.
|
|
interp = interpreter.Interpreter(
|
|
model_content=quantized_tflite_model,
|
|
experimental_op_resolver_type=interpreter.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES,
|
|
)
|
|
interp.allocate_tensors()
|
|
detail = next(
|
|
(d for d in interp.get_tensor_details() if k_dense_name in d['name'])
|
|
)
|
|
quant_params = detail['quantization_parameters']
|
|
expected_num_params = 1 if disable_per_channel_for_dense else k_num_filters
|
|
self.assertLen(quant_params['scales'], expected_num_params)
|
|
if len(quant_params['zero_points']) != 1:
|
|
self.assertLen(quant_params['zero_points'], expected_num_params)
|
|
|
|
|
|
class FromJaxModelTest(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testInvalidInputsModel(self):
|
|
if DISABLE_JAX_TEST:
|
|
return
|
|
|
|
def simple_model(input1, input2):
|
|
return jnp.sin(input1) + jnp.cos(input2)
|
|
|
|
input_tensor = jnp.zeros([10, 10])
|
|
# Invalid case: not specify serving_func
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
None, [{'input1': input_tensor}]
|
|
)
|
|
with self.assertRaisesRegex(ValueError, 'No serving func is specified.'):
|
|
converter.convert()
|
|
|
|
# Invalid case: not specify input
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[simple_model], None
|
|
)
|
|
with self.assertRaisesRegex(ValueError, 'Input tensors are not specified.'):
|
|
converter.convert()
|
|
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax([simple_model], [])
|
|
with self.assertRaisesRegex(ValueError, 'Input tensors are not specified.'):
|
|
converter.convert()
|
|
|
|
# Invalid case: not wrap input_tensor in a list.
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[simple_model], input_tensor
|
|
)
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
'The truth value of an array with more than one element is ambiguous.',
|
|
):
|
|
converter.convert()
|
|
|
|
# Invalid case: only partial inputs are provided.
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[simple_model], [[('input1', input_tensor)]]
|
|
)
|
|
with self.assertRaisesRegex(
|
|
ValueError, 'Failed to convert the given Jax function to hlo.'
|
|
):
|
|
converter.convert()
|
|
|
|
# Invalid case: serving functions length does not match input mapping.
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[simple_model, simple_model],
|
|
[[
|
|
('input1', input_tensor),
|
|
('input2', input_tensor),
|
|
]],
|
|
)
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
'Input tensor mapping len 1 does not match serving func len 2.',
|
|
):
|
|
converter.convert()
|
|
|
|
# Invalid case: multiple serving function is provided.
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[simple_model, simple_model],
|
|
[
|
|
[
|
|
('input1', input_tensor),
|
|
('input2', input_tensor),
|
|
],
|
|
[
|
|
('input1', input_tensor),
|
|
('input2', input_tensor),
|
|
],
|
|
],
|
|
)
|
|
with self.assertRaisesRegex(
|
|
ValueError, 'Currently only support single serving function.'
|
|
):
|
|
converter.convert()
|
|
|
|
@test_util.run_v2_only
|
|
def testSingleInputModel(self):
|
|
if DISABLE_JAX_TEST:
|
|
return
|
|
|
|
def single_input(input_tensor):
|
|
return jnp.sin(input_tensor)
|
|
|
|
# Convert model.
|
|
input_tensor = jnp.zeros([10, 10])
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[single_input], [[('input_tensor', input_tensor)]]
|
|
)
|
|
tflite_model = converter.convert()
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(metadata.environment.modelType, metadata_fb.ModelType.JAX)
|
|
|
|
# Check values from converted_model
|
|
input_data = np.random.random_sample((10, 10))
|
|
tf_input_data = tf.constant(input_data, dtype=np.float32)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [tf_input_data])[0]
|
|
expected_value = single_input(input_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testMultipleInputsModel(self):
|
|
if DISABLE_JAX_TEST:
|
|
return
|
|
|
|
def multiple_inputs(input1, input2):
|
|
return input1 + input2
|
|
|
|
# Convert model.
|
|
input1 = jnp.zeros([10, 10])
|
|
input2 = jnp.zeros([10, 1])
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[multiple_inputs], [[('input1', input1), ('input2', input2)]]
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted_model
|
|
input1_data = np.random.random_sample((10, 10))
|
|
tf_input1_data = tf.constant(input1_data, dtype=np.float32)
|
|
input2_data = np.random.random_sample((10, 1))
|
|
tf_input2_data = tf.constant(input2_data, dtype=np.float32)
|
|
actual_value = self._evaluateTFLiteModel(
|
|
tflite_model, [tf_input1_data, tf_input2_data]
|
|
)[0]
|
|
expected_value = multiple_inputs(input1_data, input2_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testInputSignaturesModel(self):
|
|
if DISABLE_JAX_TEST:
|
|
return
|
|
|
|
def multiple_inputs(input1, input2):
|
|
return input1 + input2
|
|
|
|
# Convert model.
|
|
input1 = jnp.zeros([10, 10])
|
|
input2 = jnp.zeros([10, 1])
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[multiple_inputs], [[('input1', input1), ('input2', input2)]]
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted_model
|
|
input1_data = np.random.random_sample((10, 10))
|
|
tf_input1_data = tf.constant(input1_data, dtype=np.float32)
|
|
input2_data = np.random.random_sample((10, 1))
|
|
tf_input2_data = tf.constant(input2_data, dtype=np.float32)
|
|
actual_value = self._evaluateTFLiteModel(
|
|
tflite_model, [tf_input1_data, tf_input2_data]
|
|
)[0]
|
|
expected_value = multiple_inputs(input1_data, input2_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testModelWithParams(self):
|
|
if DISABLE_JAX_TEST:
|
|
return
|
|
|
|
def model(inputs, weights):
|
|
return jnp.matmul(weights, inputs)
|
|
|
|
weights = np.random.random_sample((10, 10))
|
|
serving_func = functools.partial(model, weights=weights)
|
|
|
|
# Convert model
|
|
input_tensor = jnp.zeros([10, 10])
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[serving_func], [[('inputs', input_tensor)]]
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted_model
|
|
input_data = np.random.random_sample((10, 10))
|
|
tf_input_data = tf.constant(input_data, dtype=np.float32)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [tf_input_data])[0]
|
|
expected_value = serving_func(input_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testWhileLoop(self):
|
|
if DISABLE_JAX_TEST:
|
|
return
|
|
|
|
def condition(x):
|
|
return jnp.sum(x, keepdims=False) < 100
|
|
|
|
def body(x):
|
|
return jnp.add(x, 2.0)
|
|
|
|
def model(x):
|
|
result = jax.lax.while_loop(condition, body, x)
|
|
return result[0]
|
|
|
|
# Convert model.
|
|
input_tensor = jnp.zeros([3, 3])
|
|
converter = lite.TFLiteConverterV2.experimental_from_jax(
|
|
[model], [[('x', input_tensor)]]
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted_model
|
|
input_data = np.random.random_sample((3, 3))
|
|
tf_input_data = tf.constant(input_data, dtype=np.float32)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [tf_input_data])[0]
|
|
expected_value = model(input_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
|
|
class ControlFlowTest(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testCond(self):
|
|
input_data = {
|
|
'x': tf.constant([1.0, 2.0], shape=[1, 2]),
|
|
'b': tf.constant(True),
|
|
}
|
|
|
|
weights = tf.Variable([[0.1, 0.2], [0.3, 0.4]], dtype=tf.float32)
|
|
|
|
def true_fn(x):
|
|
return tf.matmul(x, weights)
|
|
|
|
def false_fn(x):
|
|
return tf.add(x, weights)
|
|
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[1, 2], dtype=tf.float32),
|
|
tf.TensorSpec(shape=(), dtype=tf.bool),
|
|
]
|
|
)
|
|
def model(x, b):
|
|
return tf.cond(
|
|
b, true_fn=lambda: true_fn(x), false_fn=lambda: false_fn(x)
|
|
)
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(**input_data)
|
|
actual_value = self._evaluateTFLiteModel(
|
|
tflite_model, [input_data['x'], input_data['b']]
|
|
)[0]
|
|
self.assertAllClose(expected_value, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testCondWithFullIntegerQuantization(self):
|
|
weights = tf.Variable([[0.1, 0.2], [0.3, 0.4]], dtype=tf.float32)
|
|
|
|
def true_fn(x):
|
|
return tf.matmul(x, weights)
|
|
|
|
def false_fn(x):
|
|
return tf.add(x, weights)
|
|
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[1, 2], dtype=tf.float32),
|
|
tf.TensorSpec(shape=(), dtype=tf.bool),
|
|
]
|
|
)
|
|
def model(x, b):
|
|
return tf.cond(
|
|
b, true_fn=lambda: true_fn(x), false_fn=lambda: false_fn(x)
|
|
)
|
|
|
|
def calibration_gen():
|
|
for _ in range(5):
|
|
yield [
|
|
np.random.uniform(-1, 1, size=(1, 2)).astype(np.float32),
|
|
tf.constant(True),
|
|
]
|
|
for _ in range(5):
|
|
yield [
|
|
np.random.uniform(-1, 1, size=(1, 2)).astype(np.float32),
|
|
tf.constant(False),
|
|
]
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = calibration_gen
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
@test_util.run_v2_only
|
|
def testConverterErrorOnControlFlowV1Ops(self):
|
|
filename = resource_loader.get_path_to_datafile(
|
|
'testdata/control_flow_v1_saved_model'
|
|
)
|
|
converter = lite.TFLiteConverterV2.from_saved_model(filename)
|
|
with self.assertRaises(convert.ConverterError) as error:
|
|
converter.convert()
|
|
self.assertIn(
|
|
'Failed to functionalize Control Flow V1 ops. Consider using Control '
|
|
'Flow V2 ops instead. See https://www.tensorflow.org/api_docs/python/'
|
|
'tf/compat/v1/enable_control_flow_v2.',
|
|
str(error.exception),
|
|
)
|
|
|
|
@test_util.run_v2_only
|
|
def testStaticRnn(self):
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((3, 10)), dtype=np.float32)
|
|
)
|
|
|
|
cell = tf.keras.layers.LSTMCell(10)
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[3, 10], dtype=tf.float32)]
|
|
)
|
|
def model(x):
|
|
seq = tf.split(x, 3, 0)
|
|
return rnn.static_rnn(cell, seq, dtype=tf.float32, sequence_length=[1])
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(input_data)[0]
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
for expected, actual in zip(expected_value, actual_value):
|
|
self.assertAllClose(expected, actual)
|
|
|
|
@test_util.run_v2_only
|
|
def testWhileLoop(self):
|
|
input_data = tf.constant([1.0, 2.0, 3.0, 4.0], shape=[2, 2])
|
|
|
|
weights = tf.Variable([[0.1, 0.2], [0.3, 0.4]], dtype=tf.float32)
|
|
|
|
def condition(x):
|
|
return tf.reduce_sum(x) < 100
|
|
|
|
def body(x):
|
|
return tf.add(x, weights)
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[2, 2], dtype=tf.float32)]
|
|
)
|
|
def model(x):
|
|
return tf.while_loop(condition, body, [x])
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(input_data)[0]
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])[0]
|
|
self.assertAllClose(expected_value, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testDynamicRnn(self):
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((3, 10, 10)), dtype=np.float32)
|
|
)
|
|
|
|
cell = tf.keras.layers.LSTMCell(10)
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[3, 10, 10], dtype=tf.float32)]
|
|
)
|
|
def model(x):
|
|
rnn_layer = tf.keras.layers.RNN([cell], return_sequences=True)
|
|
return rnn_layer(x)
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(input_data)
|
|
lite_outputs = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertLen(lite_outputs, 1)
|
|
actual_value = lite_outputs[0]
|
|
for expected, actual in zip(expected_value, actual_value):
|
|
self.assertAllClose(expected, actual)
|
|
|
|
@parameterized.named_parameters(
|
|
('LSTMBatchSizeOne', tf.keras.layers.LSTM, True),
|
|
('LSTM', tf.keras.layers.LSTM, False),
|
|
('SimpleRNNBatchSizeOne', tf.keras.layers.SimpleRNN, True),
|
|
('SimpleRNN', tf.keras.layers.SimpleRNN, False),
|
|
('GRUBatchSizeOne', tf.keras.layers.GRU, True),
|
|
('GRU', tf.keras.layers.GRU, False),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testKerasRNN(self, rnn_layer, default_to_single_batch):
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((1, 10, 10)), dtype=np.float32)
|
|
)
|
|
rnn_obj = rnn_layer(units=10, input_shape=(10, 10))
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.layers.Input(shape=(10, 10), name='input'),
|
|
rnn_obj,
|
|
])
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter._experimental_default_to_single_batch_in_tensor_list_ops = (
|
|
default_to_single_batch
|
|
)
|
|
if not default_to_single_batch:
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])[0]
|
|
|
|
# Check values from converted model.
|
|
expected_value = model.predict(input_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
@parameterized.named_parameters(
|
|
('LSTM', tf.keras.layers.LSTM),
|
|
('SimpleRNN', tf.keras.layers.SimpleRNN),
|
|
('GRU', tf.keras.layers.GRU),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testKerasRNNMultiBatches(self, rnn_layer):
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((4, 10, 10)), dtype=np.float32)
|
|
)
|
|
# Specify a fixed batch size(4) for the test model.
|
|
x = tf.keras.layers.Input(batch_shape=(4, 10, 10))
|
|
y = rnn_layer(units=10, input_shape=(10, 10))(x)
|
|
model = tf.keras.Model(inputs=[x], outputs=[y])
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
tflite_model = converter.convert()
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])[0]
|
|
|
|
# Check values from converted model.
|
|
expected_value = model.predict(input_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
@test_util.run_v2_only
|
|
def testKerasRNNLSTMFloat16Quant(self):
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((4, 10, 10)), dtype=np.float32)
|
|
)
|
|
# Specify a fixed batch size(4) for the test model.
|
|
x = tf.keras.layers.Input(batch_shape=(4, 10, 10))
|
|
y = tf.keras.layers.LSTM(units=10, input_shape=(10, 10))(x)
|
|
model = tf.keras.Model(inputs=[x], outputs=[y])
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.target_spec.supported_types = [tf.float16]
|
|
tflite_model = converter.convert()
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])[0]
|
|
|
|
# Check values from converted model.
|
|
expected_value = model.predict(input_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-03)
|
|
|
|
@parameterized.named_parameters(
|
|
('ForceToUseBatchSizeOne', True), ('DontForceToUseBatchSizeOne', False)
|
|
)
|
|
@test_util.run_v2_only
|
|
def testKerasBidirectionalRNNReturnSequence(self, default_to_single_batch):
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((1, 10, 10)), dtype=np.float32)
|
|
)
|
|
model = tf.keras.models.Sequential()
|
|
model.add(tf.keras.layers.Input(shape=(10, 10), name='input'))
|
|
model.add(
|
|
tf.keras.layers.Bidirectional(
|
|
tf.keras.layers.LSTM(units=10, return_sequences=True),
|
|
input_shape=(10, 10),
|
|
)
|
|
)
|
|
model.add(tf.keras.layers.Flatten())
|
|
model.add(tf.keras.layers.Dense(5))
|
|
model.add(tf.keras.layers.Activation('softmax'))
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter._experimental_default_to_single_batch_in_tensor_list_ops = (
|
|
default_to_single_batch
|
|
)
|
|
if not default_to_single_batch:
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])[0]
|
|
|
|
# Check values from converted model.
|
|
expected_value = model.predict(input_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
@parameterized.named_parameters(
|
|
('ForceToUseBatchSizeOne', True), ('DontForceToUseBatchSizeOne', False)
|
|
)
|
|
@test_util.run_v2_only
|
|
def testKerasBidirectionalRNN(self, default_to_single_batch):
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((1, 10, 10)), dtype=np.float32)
|
|
)
|
|
model = tf.keras.models.Sequential()
|
|
model.add(tf.keras.layers.Input(shape=(10, 10), name='input'))
|
|
model.add(tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(units=10)))
|
|
model.add(tf.keras.layers.Dense(5))
|
|
model.add(tf.keras.layers.Activation('softmax'))
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
converter._experimental_default_to_single_batch_in_tensor_list_ops = (
|
|
default_to_single_batch
|
|
)
|
|
if not default_to_single_batch:
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])[0]
|
|
|
|
# Check values from converted model.
|
|
expected_value = model.predict(input_data)
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-05)
|
|
|
|
|
|
class StridedSliceTest(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testStridedSlice(self):
|
|
input_data = tf.constant(
|
|
[
|
|
1.0,
|
|
2.0,
|
|
3.0,
|
|
4.0,
|
|
5.0,
|
|
6,
|
|
],
|
|
shape=[6],
|
|
dtype=tf.float32,
|
|
)
|
|
begin = tf.Variable([1], dtype=tf.int32)
|
|
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[6], dtype=tf.float32),
|
|
tf.TensorSpec(shape=[1], dtype=tf.int32),
|
|
]
|
|
)
|
|
def model(a, begin):
|
|
return tf.strided_slice(a, begin, begin + 3)
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(input_data, begin)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data, begin])[
|
|
0
|
|
]
|
|
self.assertAllClose(expected_value, actual_value)
|
|
|
|
|
|
class GrapplerTest(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testConstantFolding(self):
|
|
# Constant folding handles the tf.broadcast_to operation which was not
|
|
# supported by the TFLite at the time this test was added.
|
|
input_data = tf.constant(
|
|
[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], shape=[3, 3]
|
|
)
|
|
|
|
@tf.function
|
|
def func(x):
|
|
y_const = tf.constant([1.0, 2.0, 3.0])
|
|
y_broadcast = tf.broadcast_to(y_const, [3, 3])
|
|
return tf.matmul(x, y_broadcast)
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = func
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.f(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])[0]
|
|
self.assertAllClose(expected_value, actual_value)
|
|
|
|
# Enable hybrid quantization, same result
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
tflite_model = converter.convert()
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])[0]
|
|
self.assertAllClose(expected_value, actual_value)
|
|
|
|
|
|
class UnknownShapes(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testMatMul(self):
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((10, 4)), dtype=np.float32)
|
|
)
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[None, 4], dtype=tf.float32)]
|
|
)
|
|
def model(in_tensor):
|
|
shape = tf.shape(in_tensor)
|
|
fill = tf.transpose(tf.fill(shape, 1.0))
|
|
return tf.matmul(fill, in_tensor)
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(input_data)
|
|
actual_value = self._evaluateTFLiteModel(
|
|
tflite_model, [input_data], input_shapes=[([-1, 4], [10, 4])]
|
|
)[0]
|
|
self.assertAllClose(expected_value, actual_value, atol=1e-06)
|
|
|
|
def _getIntegerQuantizeModelWithUnknownShapes(self):
|
|
np.random.seed(0)
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[None, 33], dtype=tf.float32)]
|
|
)
|
|
def model(input_tensor):
|
|
"""Define a model with tf.MatMul and unknown shapes."""
|
|
# We need the tensor to have more than 1024 elements for quantize_weights
|
|
# to kick in. Thus, the [33, 33] shape.
|
|
const_tensor = tf.constant(
|
|
np.random.uniform(low=-10.0, high=10.0, size=[33, 33]),
|
|
shape=[33, 33],
|
|
dtype=tf.float32,
|
|
name='inputB',
|
|
)
|
|
|
|
shape = tf.shape(input_tensor)
|
|
fill = tf.transpose(tf.fill(shape, 1.0))
|
|
mult = tf.matmul(fill, input_tensor)
|
|
return tf.matmul(mult, const_tensor)
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = model
|
|
concrete_func = root.f.get_concrete_function()
|
|
|
|
def calibration_gen():
|
|
for batch in range(5, 20, 5):
|
|
for _ in range(5):
|
|
yield [np.random.uniform(-1, 1, size=(batch, 33)).astype(np.float32)]
|
|
|
|
return root, concrete_func, calibration_gen
|
|
|
|
@test_util.run_v2_only
|
|
def testMatMulQuantize(self):
|
|
root, concrete_func, _ = self._getIntegerQuantizeModelWithUnknownShapes()
|
|
float_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
float_tflite_model = float_converter.convert()
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
|
|
# The default input and output types should be float.
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual(np.float32, input_details[0]['dtype'])
|
|
self.assertAllEqual([-1, 33], input_details[0]['shape_signature'])
|
|
|
|
# Ensure that the quantized weights tflite model is smaller.
|
|
self.assertLess(len(quantized_tflite_model), len(float_tflite_model))
|
|
|
|
@test_util.run_v2_only
|
|
def testMatMulCalibrateAndQuantize(self):
|
|
root, concrete_func, calibration_gen = (
|
|
self._getIntegerQuantizeModelWithUnknownShapes()
|
|
)
|
|
float_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
float_tflite_model = float_converter.convert()
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
|
|
# The default input and output types should be float.
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
self.assertLen(input_details, 1)
|
|
self.assertEqual(np.float32, input_details[0]['dtype'])
|
|
self.assertAllEqual([-1, 33], input_details[0]['shape_signature'])
|
|
|
|
# Ensure that the quantized weights tflite model is smaller.
|
|
self.assertLess(len(quantized_tflite_model), len(float_tflite_model))
|
|
|
|
def testBatchMatMul(self):
|
|
input_data_1 = tf.constant(
|
|
np.array(np.random.random_sample((1, 256, 256)), dtype=np.float32)
|
|
)
|
|
input_data_2 = tf.constant(
|
|
np.array(np.random.random_sample((1, 256, 256)), dtype=np.float32)
|
|
)
|
|
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[None, 256, 256], dtype=tf.float32),
|
|
tf.TensorSpec(shape=[None, 256, 256], dtype=tf.float32),
|
|
]
|
|
)
|
|
def model(in_tensor_1, in_tensor_2):
|
|
return tf.matmul(in_tensor_1, in_tensor_2)
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(input_data_1, input_data_2)
|
|
actual_value = self._evaluateTFLiteModel(
|
|
tflite_model,
|
|
[input_data_1, input_data_2],
|
|
input_shapes=[([-1, 256, 256], [1, 256, 256])],
|
|
)[0]
|
|
self.assertAllClose(expected_value, actual_value, atol=4)
|
|
|
|
def testBatchMatMulInputInt8Int8OutputInt32(self):
|
|
input_data_1 = tf.constant(
|
|
np.array(
|
|
np.random.random_integers(-128, high=127, size=(1, 20, 30)),
|
|
dtype=np.int8,
|
|
)
|
|
)
|
|
input_data_2 = tf.constant(
|
|
np.array(
|
|
np.random.random_integers(-128, high=127, size=(1, 30, 10)),
|
|
dtype=np.int8,
|
|
)
|
|
)
|
|
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[None, 20, 30], dtype=tf.int8),
|
|
tf.TensorSpec(shape=[None, 30, 10], dtype=tf.int8),
|
|
]
|
|
)
|
|
def model(in_tensor_1, in_tensor_2):
|
|
return tf.matmul(in_tensor_1, in_tensor_2, output_type=tf.int32)
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(input_data_1, input_data_2)
|
|
actual_value = self._evaluateTFLiteModel(
|
|
tflite_model,
|
|
[input_data_1, input_data_2],
|
|
input_shapes=[([-1, 20, 30], [1, 20, 30]), ([-1, 30, 10], [1, 30, 10])],
|
|
)[0]
|
|
self.assertAllEqual(expected_value, actual_value)
|
|
|
|
def testBatchMatMulHybrid(self):
|
|
# Test model that does batch matmul of:
|
|
# lhs input (1, 256, 128), rhs const (1, 128, 256).
|
|
# For dynamic range quantization situation, this will result in hybrid batch
|
|
# matmul, where lhs type is float32 and rhs type is int8.
|
|
|
|
# Intentionally set lhs, rhs sizes to satisfy following conditions:
|
|
# 1. rhs const num_elements >= 1024, since dynamic range quantization
|
|
# requires const tensor num_elements to be larger than
|
|
# min_elements_for_weights (which defaults to 1024).
|
|
# (https://github.com/tensorflow/tensorflow/blob/25e649ac3688655547da998eba2715cf70b3e5c9/tensorflow/compiler/mlir/lite/transforms/prepare_quantize_dynamic_range.cc#L262)
|
|
# 2. batch_size (256) > accum_dim_size (128) and
|
|
# num_units (256) > accum_dim_size (128), to test if the sizes are set
|
|
# correctly according to dimensions. See HybridAsymmetricBatchMatMulOpTest
|
|
# tests in
|
|
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/batch_matmul_test.cc.
|
|
input_data = tf.constant(
|
|
np.array(np.random.random_sample((1, 256, 128)), dtype=np.float32)
|
|
)
|
|
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[None, 256, 128], dtype=tf.float32)
|
|
]
|
|
)
|
|
def model(in_tensor):
|
|
rhs = tf.constant(
|
|
np.array(np.random.random_sample((1, 128, 256)), dtype=np.float32)
|
|
)
|
|
return tf.matmul(in_tensor, rhs)
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = concrete_func(input_data)
|
|
actual_value = self._evaluateTFLiteModel(
|
|
tflite_model,
|
|
[input_data],
|
|
input_shapes=[([-1, 256, 128], [1, 256, 128])],
|
|
)[0]
|
|
self.assertAllClose(expected_value, actual_value, atol=4)
|
|
|
|
def testSizeInvalid(self):
|
|
@tf.function(
|
|
input_signature=[
|
|
tf.TensorSpec(shape=[1, None, 16, 3], dtype=tf.float32)
|
|
]
|
|
)
|
|
def model(in_tensor):
|
|
return in_tensor + in_tensor
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
# Test invalid shape. None after 1st dimension. Run with TOCO in order to
|
|
# invoke shape checking code.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
converter.experimental_new_converter = False
|
|
with self.assertRaises(ValueError) as error:
|
|
converter.convert()
|
|
self.assertEqual(
|
|
'None is only supported in the 1st dimension. Tensor '
|
|
"'in_tensor' has invalid shape '[1, None, 16, 3]'.",
|
|
str(error.exception),
|
|
)
|
|
|
|
|
|
class ResourceAndVariantTypes(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testVariants(self):
|
|
@tf.function(input_signature=[tf.TensorSpec(shape=[1], dtype=tf.float32)])
|
|
def model(v):
|
|
m = map_ops.empty_tensor_map()
|
|
k = tf.constant(1.0)
|
|
p = tf.add(k, v)
|
|
with ops.control_dependencies([m]):
|
|
m2 = map_ops.tensor_map_insert(m, p, v)
|
|
with ops.control_dependencies([m2]):
|
|
return map_ops.tensor_map_size(m2)
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([1.0], dtype=np.float32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(1, actual_value)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(1, actual_value)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(1, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testVariantsWithCond(self):
|
|
def create_v1_saved_model():
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'variants_with_cond')
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
m = map_ops.empty_tensor_map()
|
|
|
|
def body(i, m):
|
|
m = map_ops.tensor_map_insert(m, i, i)
|
|
return i + 1, m
|
|
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1], dtype=tf.int32, name='input'
|
|
)
|
|
_, result_m = tf.cond(
|
|
in_tensor < 10,
|
|
lambda: body(in_tensor, m),
|
|
lambda: body(in_tensor + 1, m),
|
|
)
|
|
out_tensor = in_tensor + map_ops.tensor_map_size(result_m)
|
|
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'z': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
saved_model_dir = create_v1_saved_model()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([0], dtype=np.int32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
expected_value = np.array([1], dtype=np.int32)
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(expected_value, actual_value)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(expected_value, actual_value)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(expected_value, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testVariantsWithWhile(self):
|
|
def create_v1_saved_model():
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'variants_with_while')
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
m = map_ops.empty_tensor_map()
|
|
|
|
def cond(i, m):
|
|
del m
|
|
return i < 10
|
|
|
|
def body(i, m):
|
|
m = map_ops.tensor_map_insert(m, i, i)
|
|
return i + 1, m
|
|
|
|
_, result_m = tf.while_loop(cond, body, [0, m])
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1], dtype=tf.int32, name='input'
|
|
)
|
|
out_tensor = in_tensor + map_ops.tensor_map_size(result_m)
|
|
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'z': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
saved_model_dir = create_v1_saved_model()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([0], dtype=np.int32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(10, actual_value)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(10, actual_value)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(10, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testResources(self):
|
|
def create_v1_saved_model():
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_resources')
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1], dtype=tf.float32, name='input'
|
|
)
|
|
|
|
stack = tf.raw_ops.StackV2(max_size=10, elem_type=tf.float32)
|
|
w = tf.raw_ops.StackPushV2(handle=stack, elem=in_tensor)
|
|
with ops.control_dependencies([w]):
|
|
a = in_tensor + in_tensor
|
|
with ops.control_dependencies([a]):
|
|
out_tensor = a + tf.raw_ops.StackPopV2(
|
|
handle=stack, elem_type=tf.float32
|
|
)
|
|
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'z': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
saved_model_dir = create_v1_saved_model()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([1.0], dtype=np.float32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(3.0, actual_value)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(3.0, actual_value)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(3.0, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testResourcesWithCond(self):
|
|
def create_v1_saved_model():
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'resources_with_cond')
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1], dtype=tf.float32, name='input'
|
|
)
|
|
|
|
def body(i, arr):
|
|
n = tf.raw_ops.StackPushV2(
|
|
handle=arr, elem=tf.cast(i, dtype=tf.float32)
|
|
)
|
|
return n, arr
|
|
|
|
arr = tf.raw_ops.StackV2(max_size=10, elem_type=tf.float32)
|
|
n, result_arr = tf.cond(
|
|
in_tensor < 10, lambda: body(0, arr), lambda: body(1, arr)
|
|
)
|
|
|
|
with ops.control_dependencies([result_arr, n]):
|
|
out_tensor = tf.raw_ops.StackPopV2(
|
|
handle=result_arr, elem_type=tf.float32
|
|
)
|
|
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'a': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
saved_model_dir = create_v1_saved_model()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([1.0], dtype=np.float32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(0.0, actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testResourcesWithWhile(self):
|
|
def create_v1_saved_model():
|
|
saved_model_dir = os.path.join(
|
|
self.get_temp_dir(), 'resources_with_while'
|
|
)
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1], dtype=tf.float32, name='input'
|
|
)
|
|
|
|
def cond(i, arr, m):
|
|
del arr
|
|
del m
|
|
return i < 10
|
|
|
|
def body(i, arr, m):
|
|
del m
|
|
n = tf.raw_ops.StackPushV2(
|
|
handle=arr, elem=tf.cast(i, dtype=tf.float32)
|
|
)
|
|
return i + 1, arr, n
|
|
|
|
arr = tf.raw_ops.StackV2(max_size=10, elem_type=tf.float32)
|
|
_, result_arr, n = tf.while_loop(cond, body, [0, arr, 0.0])
|
|
|
|
with ops.control_dependencies([result_arr, n]):
|
|
out_tensor = tf.raw_ops.StackPopV2(
|
|
handle=result_arr, elem_type=tf.float32
|
|
)
|
|
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'a': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
saved_model_dir = create_v1_saved_model()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([1.0], dtype=np.float32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(9.0, actual_value)
|
|
|
|
@parameterized.named_parameters(
|
|
('EnableLoweringTensorListOps', True),
|
|
('DisableLoweringTensorListOps', False),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testTensorListWithStaticSize(self, lower_tensor_list_ops):
|
|
def create_v1_saved_model():
|
|
saved_model_dir = os.path.join(
|
|
self.get_temp_dir(), 'simple_mutable_variable'
|
|
)
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1], dtype=tf.float32, name='input'
|
|
)
|
|
|
|
ta = tf.TensorArray(
|
|
tf.float32, size=3, dynamic_size=False, clear_after_read=False
|
|
)
|
|
ta = ta.write(0, 10.0)
|
|
ta = ta.write(1, 20.0)
|
|
ta = ta.write(2, 30.0)
|
|
|
|
out_tensor = ta.read(0) + ta.read(2)
|
|
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'z': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
saved_model_dir = create_v1_saved_model()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
if not lower_tensor_list_ops:
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
converter._experimental_lower_tensor_list_ops = lower_tensor_list_ops
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([1.0], dtype=np.float32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(40.0, actual_value)
|
|
|
|
@parameterized.named_parameters(
|
|
('EnableLoweringTensorListOps', True),
|
|
('DisableLoweringTensorListOps', False),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testTensorListWithDynamicSize(self, lower_tensor_list_ops):
|
|
def create_v1_saved_model():
|
|
saved_model_dir = os.path.join(
|
|
self.get_temp_dir(), 'simple_mutable_variable'
|
|
)
|
|
with tf.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1], dtype=tf.float32, name='input'
|
|
)
|
|
|
|
ta = tf.TensorArray(
|
|
tf.float32, size=0, dynamic_size=True, clear_after_read=False
|
|
)
|
|
ta = ta.write(0, 10.0)
|
|
ta = ta.write(1, 20.0)
|
|
ta = ta.write(2, 30.0)
|
|
|
|
out_tensor = ta.read(0) + ta.read(2)
|
|
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'z': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
return saved_model_dir
|
|
|
|
saved_model_dir = create_v1_saved_model()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
if lower_tensor_list_ops:
|
|
with self.assertRaises(convert.ConverterError) as error:
|
|
converter.convert()
|
|
self.assertIn(
|
|
'Lowering tensor list ops is failed. Please consider using Select '
|
|
'TF ops and disabling `_experimental_lower_tensor_list_ops` flag in '
|
|
'the TFLite converter object.',
|
|
str(error.exception),
|
|
)
|
|
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
input_data = np.array([1.0], dtype=np.float32)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(40.0, actual_value)
|
|
|
|
|
|
class CalibrateAndQuantizeWithCustomOpTest(lite_v2_test_util.ModelTest):
|
|
|
|
def _createGraphWithCustomOp(self):
|
|
# Create a graph that has one double op.
|
|
np.random.seed(0)
|
|
|
|
saved_model_dir = os.path.join(self.get_temp_dir(), 'double_model')
|
|
with ops.Graph().as_default():
|
|
with tf.compat.v1.Session() as sess:
|
|
in_tensor = tf.compat.v1.placeholder(
|
|
shape=[1, 4], dtype=dtypes.float32, name='input'
|
|
)
|
|
out_tensor = double_op.double(in_tensor)
|
|
inputs = {'x': in_tensor}
|
|
outputs = {'z': out_tensor}
|
|
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
|
|
|
|
def calibration_gen():
|
|
for _ in range(100):
|
|
yield [np.random.uniform(-1, 1, size=(1, 4)).astype(np.float32)]
|
|
|
|
return (saved_model_dir, calibration_gen)
|
|
|
|
def testCustomOpRegistererByName(self):
|
|
"""Test a calibration with custom op registered by name."""
|
|
saved_model_dir, calibration_gen = self._createGraphWithCustomOp()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = calibration_gen
|
|
converter.allow_custom_ops = True
|
|
converter.target_spec._experimental_custom_op_registerers = [
|
|
'TF_TestRegisterer'
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
self.assertGreater(test_registerer.get_num_test_registerer_calls(), 0)
|
|
self.assertIn('Double', tflite_test_util.get_ops_list(tflite_model))
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(metadata.options.allowCustomOps, True)
|
|
|
|
# Check the model works with custom ops.
|
|
interp = interpreter.InterpreterWithCustomOps(
|
|
model_content=tflite_model, custom_op_registerers=['TF_TestRegisterer']
|
|
)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
test_input = np.array([[0.0, 0.1, 0.2, 0.3]], dtype=np.float32)
|
|
interp.set_tensor(input_details[0]['index'], test_input)
|
|
interp.invoke()
|
|
|
|
output_details = interp.get_output_details()
|
|
expected_output = np.array([[0.0, 0.2, 0.4, 0.6]], dtype=np.float32)
|
|
output_data = interp.get_tensor(output_details[0]['index'])
|
|
self.assertArrayNear(expected_output[0], output_data[0], err=1e-2)
|
|
|
|
def testCustomOpRegistererByFunc(self):
|
|
"""Test a calibration with custom op registered by function."""
|
|
saved_model_dir, calibration_gen = self._createGraphWithCustomOp()
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = calibration_gen
|
|
converter.allow_custom_ops = True
|
|
converter.target_spec._experimental_custom_op_registerers = [
|
|
test_registerer.TF_TestRegisterer
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertTrue(tflite_model)
|
|
self.assertGreater(test_registerer.get_num_test_registerer_calls(), 0)
|
|
self.assertIn('Double', tflite_test_util.get_ops_list(tflite_model))
|
|
|
|
# Check the model works with custom ops.
|
|
interp = interpreter.InterpreterWithCustomOps(
|
|
model_content=tflite_model,
|
|
custom_op_registerers=[test_registerer.TF_TestRegisterer],
|
|
)
|
|
interp.allocate_tensors()
|
|
input_details = interp.get_input_details()
|
|
test_input = np.array([[0.0, 0.1, 0.2, 0.3]], dtype=np.float32)
|
|
interp.set_tensor(input_details[0]['index'], test_input)
|
|
interp.invoke()
|
|
|
|
output_details = interp.get_output_details()
|
|
expected_output = np.array([[0.0, 0.2, 0.4, 0.6]], dtype=np.float32)
|
|
output_data = interp.get_tensor(output_details[0]['index'])
|
|
self.assertArrayNear(expected_output[0], output_data[0], err=1e-2)
|
|
|
|
def testCustomOpRegistererFailure(self):
|
|
"""Test a calibration with wrong custom op registerer."""
|
|
saved_model_dir, calibration_gen = self._createGraphWithCustomOp()
|
|
|
|
bogus_name = 'CompletelyBogusRegistererName'
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
converter.representative_dataset = calibration_gen
|
|
converter.allow_custom_ops = True
|
|
converter.target_spec._experimental_custom_op_registerers = [bogus_name]
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "Looking up symbol '" + bogus_name + "' failed"
|
|
):
|
|
converter.convert()
|
|
|
|
|
|
class IntermediatesTest(lite_v2_test_util.ModelTest):
|
|
|
|
def _run(self, experimental_preserve_all_tensors):
|
|
@tf.function
|
|
def f(x):
|
|
y = tf.add(x, x, name='y')
|
|
z = tf.add(y, y, name='z')
|
|
w = tf.add(z, z, name='w')
|
|
return w
|
|
|
|
# NOTE this is exactly representable as a float as are the intermediates of
|
|
# f. So direct comparison is ok below.
|
|
|
|
input_data = np.array(2.0, np.float32)
|
|
concrete_func = f.get_concrete_function(input_data)
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], f
|
|
)
|
|
tflite_model = converter.convert()
|
|
interp = interpreter.Interpreter(
|
|
model_content=tflite_model,
|
|
experimental_preserve_all_tensors=experimental_preserve_all_tensors,
|
|
)
|
|
interp.allocate_tensors()
|
|
interp.set_tensor(interp.get_input_details()[0]['index'], input_data)
|
|
interp.invoke()
|
|
out = interp.get_tensor(interp.get_output_details()[0]['index'])
|
|
tensors = {}
|
|
for t in interp.get_tensor_details():
|
|
# With Tensorflow Lite default delegate applied to the model graph, the
|
|
# access to original tensors of a delegated op could cause a ValueError
|
|
# (i.e. 'Tensor data is null. Run allocate_tensors() first') to be thrown
|
|
# out because the tensor memory isn't allocated at all.
|
|
val = None
|
|
try:
|
|
val = interp.get_tensor(t['index'])
|
|
except ValueError:
|
|
pass
|
|
tensors.update({t['name']: val})
|
|
return (tensors, out)
|
|
|
|
def testPreserve(self):
|
|
tensors, result = self._run(experimental_preserve_all_tensors=True)
|
|
# All intermediates should be true and result be true.
|
|
self.assertAllClose(tensors['x'], 2.0)
|
|
self.assertAllClose(tensors['y'], 4.0)
|
|
self.assertAllClose(tensors['z'], 8.0)
|
|
self.assertAllClose(result, 16.0)
|
|
|
|
def testNoPreserve(self):
|
|
tensors, result = self._run(experimental_preserve_all_tensors=False)
|
|
# One of them should be wrong if preserve is not true, but result should be
|
|
# ok. Input should still be ok for repeated invocation.
|
|
self.assertAllClose(tensors['x'], 2.0)
|
|
self.assertTrue(tensors['y'] != 4.0 or tensors['z'] != 8.0)
|
|
self.assertAllClose(result, 16.0)
|
|
|
|
|
|
class DatasetOpsTest(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testReduceDataset(self):
|
|
@tf.function
|
|
def model():
|
|
dataset = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4])
|
|
output = dataset.reduce(np.int32(0), lambda x, y: x + y)
|
|
return output
|
|
|
|
concrete_func = model.get_concrete_function()
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], model
|
|
)
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS,
|
|
lite.OpsSet.SELECT_TF_OPS,
|
|
]
|
|
tflite_model = converter.convert()
|
|
self.assertIsNotNone(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
output_details = interp.get_output_details()
|
|
|
|
interp.allocate_tensors()
|
|
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertEqual(10, actual_value)
|
|
|
|
|
|
class SparsityTest(lite_v2_test_util.ModelTest):
|
|
|
|
def _getSparsificableModel(self, matrix_b_values):
|
|
np.random.seed(0)
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@tf.function(
|
|
input_signature=[tf.TensorSpec(shape=[16, 4], dtype=tf.float32)]
|
|
)
|
|
def func(inp):
|
|
matrix_b = tf.constant(matrix_b_values, dtype=tf.float32)
|
|
matrix_b = tf.reshape(matrix_b, [4, 8])
|
|
matmul = tf.matmul(inp, matrix_b, transpose_a=False, transpose_b=False)
|
|
output = tf.nn.relu(matmul, name='output')
|
|
return output
|
|
|
|
root.f = func
|
|
to_save = root.f.get_concrete_function()
|
|
return (root, to_save)
|
|
|
|
def testRandomSparsity(self):
|
|
# pyformat: disable
|
|
matrix_b_values = [
|
|
0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 0, 1,
|
|
]
|
|
# pyformat: enable
|
|
root, func = self._getSparsificableModel(matrix_b_values)
|
|
float_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
float_converter.optimizations = [lite.Optimize.EXPERIMENTAL_SPARSITY]
|
|
float_tflite_model = float_converter.convert()
|
|
self.assertIsNotNone(float_tflite_model)
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(float_tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertAllEqual(
|
|
[metadata_fb.ModelOptimizationMode.RANDOM_SPARSITY],
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
|
|
def testBlockSparsity(self):
|
|
# pyformat: disable
|
|
matrix_b_values = [
|
|
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
|
0, 0, 0, 0, 0, 0, 1, 0,
|
|
]
|
|
# pyformat: enable
|
|
root, func = self._getSparsificableModel(matrix_b_values)
|
|
float_converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[func], root
|
|
)
|
|
float_converter.optimizations = [lite.Optimize.EXPERIMENTAL_SPARSITY]
|
|
float_tflite_model = float_converter.convert()
|
|
self.assertIsNotNone(float_tflite_model)
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(float_tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertAllEqual(
|
|
[metadata_fb.ModelOptimizationMode.BLOCK_SPARSITY],
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
|
|
@parameterized.named_parameters(
|
|
('_PerChannelQuantForDense', False),
|
|
('_PerTensorQuantForDense', True),
|
|
)
|
|
def testQuantizedBlockSparsity(
|
|
self, disable_per_channel_quantization_for_dense_layers
|
|
):
|
|
weight_values = np.array([
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 2, 0, 0, 0, 0, 5, 0, 0, 0, 3, 0, 0, 0, 1, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[3, 0, 7, 0, 0, 0, -6, -2, 0, 0, 0, 0, 0, -2, 0, 6],
|
|
])
|
|
|
|
custom_init = tf.constant_initializer(weight_values.transpose())
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.layers.Dense(
|
|
8, kernel_initializer=custom_init, input_shape=[16]
|
|
)
|
|
])
|
|
|
|
def calibration_gen():
|
|
for _ in range(10):
|
|
yield [np.random.uniform(-1, 1, size=(1, 16)).astype(np.float32) * 16]
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
quantized_converter.optimizations = [
|
|
lite.Optimize.EXPERIMENTAL_SPARSITY,
|
|
lite.Optimize.DEFAULT,
|
|
]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
quantized_converter._experimental_disable_per_channel_quantization_for_dense_layers = (
|
|
disable_per_channel_quantization_for_dense_layers
|
|
)
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(quantized_tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(
|
|
metadata.environment.tensorflowVersion.decode('utf-8'),
|
|
versions.__version__,
|
|
)
|
|
self.assertEqual(metadata.environment.apiVersion, 2)
|
|
self.assertAllEqual(
|
|
[
|
|
metadata_fb.ModelOptimizationMode.PTQ_FULL_INTEGER,
|
|
metadata_fb.ModelOptimizationMode.BLOCK_SPARSITY,
|
|
],
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
interp.allocate_tensors()
|
|
input_data = np.array(
|
|
[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]],
|
|
dtype=np.float32,
|
|
)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertArrayNear(
|
|
np.array([0, 87, 0, 0, 0, 0, 0, 34], dtype=np.float32),
|
|
actual_value.flatten(),
|
|
err=1,
|
|
)
|
|
|
|
def testQuantizedButNotEnoughBlockSparsity(self):
|
|
# Sparsity level is 25%, which is not enough to apply the sparse conversion.
|
|
weight_values = np.array([
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[4, 4, -3, 4, 4, 1, -2, -2, 1, 3, 4, 1, 1, 1, -4, -5],
|
|
[1, 1, 5, -1, 3, -1, 1, -3, 4, -3, 2, -3, 3, -1, 3, -4],
|
|
[0, -3, -2, 5, 4, 2, 1, 4, -4, 4, 1, -2, 3, -2, -2, -1],
|
|
])
|
|
|
|
custom_init = tf.constant_initializer(weight_values.transpose())
|
|
model = tf.keras.models.Sequential([
|
|
tf.keras.layers.Dense(
|
|
4, kernel_initializer=custom_init, input_shape=[16]
|
|
)
|
|
])
|
|
|
|
def calibration_gen():
|
|
for _ in range(10):
|
|
yield [np.random.uniform(-1, 1, size=(1, 16)).astype(np.float32) * 16]
|
|
|
|
quantized_converter = lite.TFLiteConverterV2.from_keras_model(model)
|
|
quantized_converter.optimizations = [
|
|
lite.Optimize.EXPERIMENTAL_SPARSITY,
|
|
lite.Optimize.DEFAULT,
|
|
]
|
|
quantized_converter.representative_dataset = calibration_gen
|
|
quantized_tflite_model = quantized_converter.convert()
|
|
self.assertIsNotNone(quantized_tflite_model)
|
|
|
|
# Check the conversion metadata.
|
|
metadata = util.get_conversion_metadata(quantized_tflite_model)
|
|
self.assertIsNotNone(metadata)
|
|
self.assertEqual(
|
|
metadata.environment.tensorflowVersion.decode('utf-8'),
|
|
versions.__version__,
|
|
)
|
|
self.assertEqual(metadata.environment.apiVersion, 2)
|
|
self.assertAllEqual(
|
|
[
|
|
metadata_fb.ModelOptimizationMode.PTQ_FULL_INTEGER,
|
|
],
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
self.assertNotIn(
|
|
metadata_fb.ModelOptimizationMode.RANDOM_SPARSITY,
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
self.assertNotIn(
|
|
metadata_fb.ModelOptimizationMode.BLOCK_SPARSITY,
|
|
metadata.options.modelOptimizationModes,
|
|
)
|
|
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=quantized_tflite_model)
|
|
input_details = interp.get_input_details()
|
|
output_details = interp.get_output_details()
|
|
interp.allocate_tensors()
|
|
input_data = np.array(
|
|
[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]],
|
|
dtype=np.float32,
|
|
)
|
|
interp.set_tensor(input_details[0]['index'], input_data)
|
|
interp.invoke()
|
|
actual_value = interp.get_tensor(output_details[0]['index'])
|
|
self.assertArrayNear(
|
|
np.array([0, -3, 4, 35], dtype=np.float32),
|
|
actual_value.flatten(),
|
|
err=1,
|
|
)
|
|
|
|
|
|
class BufferOffsetTest(lite_v2_test_util.ModelTest):
|
|
|
|
@test_util.run_v2_only
|
|
def testCOncreteFunctionFloat(self):
|
|
root = self._getSimpleVariableModel()
|
|
input_data = tf.constant(1.0, shape=[1])
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
converter._experimental_use_buffer_offset = True
|
|
tflite_model = converter.convert()
|
|
|
|
# Check output value from converted model.
|
|
expected_value = root.f(input_data)
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value.numpy(), actual_value)
|
|
|
|
@test_util.run_v2_only
|
|
def testConcreteFunctionStringInput(self):
|
|
class Model(tf.Module):
|
|
|
|
@tf.function
|
|
def __call__(self, x):
|
|
return x
|
|
|
|
root = Model()
|
|
concrete_func = root.__call__.get_concrete_function(
|
|
tf.constant([str(x) for x in range(11)])
|
|
)
|
|
# Convert model.
|
|
converter = lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
converter._experimental_use_buffer_offset = True
|
|
tflite_model = converter.convert()
|
|
input_data = tf.constant(
|
|
[str(x) for x in range(11)], shape=(11,), dtype=tf.dtypes.string
|
|
)
|
|
# Check values from converted model.
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
interp.allocate_tensors()
|
|
my_signature = interp.get_signature_runner()
|
|
|
|
with self.assertRaises(ValueError) as error:
|
|
_ = my_signature(x=input_data)
|
|
self.assertIn(
|
|
'Passed in value type is not a numpy array, got type ',
|
|
str(error.exception),
|
|
)
|
|
|
|
@test_util.run_v2_only
|
|
def testSavedModelSignatureDefs(self):
|
|
"""Test converting SignatureDef is correct and uses SignatureDef API."""
|
|
root = self._getMultiFunctionModel()
|
|
input_data_0 = tf.constant(1.0, shape=[1])
|
|
input_data_1 = tf.constant(3.0, shape=[1])
|
|
mul_add_func = root.mul_add.get_concrete_function(
|
|
input_data_1, input_data_0
|
|
)
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
save.save(root, save_dir, {'mul_add': mul_add_func})
|
|
|
|
converter = lite.TFLiteConverterV2.from_saved_model(
|
|
save_dir, signature_keys=['mul_add']
|
|
)
|
|
converter._experimental_use_buffer_offset = True
|
|
tflite_model = converter.convert()
|
|
|
|
# Check values from converted model.
|
|
expected_value = root.mul_add(input_data_1, input_data_0)
|
|
interp = interpreter.Interpreter(model_content=tflite_model)
|
|
signature_defs = interp.get_signature_list()
|
|
results = self._evaluateTFLiteModelUsingSignatureDef(
|
|
tflite_model, 'mul_add', {'y': input_data_0, 'x': input_data_1}
|
|
)
|
|
self.assertEqual(list(results.keys()), ['output_0'])
|
|
self.assertEqual(expected_value.numpy(), results['output_0'])
|
|
|
|
# Verify the SignatureDef structure returned is as expected.
|
|
self.assertLen(signature_defs, 1)
|
|
self.assertEqual(list(signature_defs.keys()), ['mul_add'])
|
|
self.assertLen(signature_defs.values(), 1)
|
|
self.assertEqual(
|
|
list(signature_defs['mul_add'].keys()), ['inputs', 'outputs']
|
|
)
|
|
self.assertCountEqual(signature_defs['mul_add']['inputs'], ['x', 'y'])
|
|
self.assertEqual(list(signature_defs['mul_add']['outputs']), ['output_0'])
|
|
|
|
|
|
class BoundaryValueTest(lite_v2_test_util.ModelTest):
|
|
|
|
@parameterized.named_parameters(
|
|
('EnableCanonicalizeInfAsMaxMinFloatFromSavedModel', True, True),
|
|
('DisableCanonicalizeInfAsMaxMinFloatFromSavedModel', False, True),
|
|
('EnableCanonicalizeInfAsMaxMinFloatFromConcreteFunc', True, False),
|
|
('DisableCanonicalizeInfAsMaxMinFloatFromConcreteFunc', False, False),
|
|
)
|
|
@test_util.run_v2_only
|
|
def testFloatBoundaryValue(self, is_canonicalized, is_from_saved_model):
|
|
root = self._getInfFloatModel()
|
|
input_data = None
|
|
concrete_func = root.f.get_concrete_function(input_data)
|
|
|
|
mdl = tf.Module()
|
|
mdl.f = concrete_func
|
|
|
|
def _get_converter() -> lite.TFLiteConverterV2:
|
|
if is_from_saved_model:
|
|
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
|
|
tf.saved_model.save(mdl, save_dir)
|
|
return lite.TFLiteConverterV2.from_saved_model(save_dir)
|
|
return lite.TFLiteConverterV2.from_concrete_functions(
|
|
[concrete_func], root
|
|
)
|
|
|
|
converter = _get_converter()
|
|
converter.canonicalizing_inf_as_min_max_float = is_canonicalized
|
|
tflite_model = converter.convert()
|
|
|
|
# Check output value from converted model.
|
|
expected_value = [np.finfo(np.float32).max if is_canonicalized else np.inf]
|
|
actual_value = self._evaluateTFLiteModel(tflite_model, [input_data])
|
|
self.assertEqual(expected_value, actual_value)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test.main()
|