From fbc49991ab6b7616f984735f488320251ddc4402 Mon Sep 17 00:00:00 2001 From: Laura Pak Date: Mon, 28 Oct 2024 10:43:16 -0700 Subject: [PATCH] Move tensorflow lite python calls to ai-edge-litert. PiperOrigin-RevId: 690670455 --- mediapipe/model_maker/python/core/utils/BUILD | 2 -- mediapipe/model_maker/python/core/utils/model_util.py | 3 +-- .../model_maker/python/core/utils/quantization_test.py | 5 ++--- mediapipe/model_maker/requirements.txt | 1 - mediapipe/model_maker/requirements_bazel.txt | 1 - mediapipe/model_maker/requirements_lock.txt | 7 +------ 6 files changed, 4 insertions(+), 15 deletions(-) diff --git a/mediapipe/model_maker/python/core/utils/BUILD b/mediapipe/model_maker/python/core/utils/BUILD index ef6e652845..4b8ff1bf97 100644 --- a/mediapipe/model_maker/python/core/utils/BUILD +++ b/mediapipe/model_maker/python/core/utils/BUILD @@ -44,7 +44,6 @@ py_library( deps = [ ":quantization", "//mediapipe/model_maker/python/core/data:dataset", - "@model_maker_pip_deps_ai_edge_litert_nightly//:pkg", "@model_maker_pip_deps_numpy//:pkg", "@model_maker_pip_deps_tensorflow//:pkg", ], @@ -173,7 +172,6 @@ py_test( ":quantization", ":test_util", "@model_maker_pip_deps_absl_py//:pkg", - "@model_maker_pip_deps_ai_edge_litert_nightly//:pkg", "@model_maker_pip_deps_tensorflow//:pkg", ], ) diff --git a/mediapipe/model_maker/python/core/utils/model_util.py b/mediapipe/model_maker/python/core/utils/model_util.py index 0a13095e6a..32b509797f 100644 --- a/mediapipe/model_maker/python/core/utils/model_util.py +++ b/mediapipe/model_maker/python/core/utils/model_util.py @@ -28,7 +28,6 @@ from mediapipe.model_maker.python.core.data import dataset from mediapipe.model_maker.python.core.utils import quantization -from ai_edge_litert import interpreter as tfl_interpreter DEFAULT_SCALE, DEFAULT_ZERO_POINT = 0, 0 ESTIMITED_STEPS_PER_EPOCH = 1000 @@ -274,7 +273,7 @@ def __init__(self, tflite_model: bytearray): Args: tflite_model: A valid flatbuffer representing the TFLite model. """ - self.interpreter = tfl_interpreter.Interpreter(model_content=tflite_model) + self.interpreter = tf.lite.Interpreter(model_content=tflite_model) self.interpreter.allocate_tensors() self.input_details = self.interpreter.get_input_details() self.output_details = self.interpreter.get_output_details() diff --git a/mediapipe/model_maker/python/core/utils/quantization_test.py b/mediapipe/model_maker/python/core/utils/quantization_test.py index 0164d39bf8..57523d4056 100644 --- a/mediapipe/model_maker/python/core/utils/quantization_test.py +++ b/mediapipe/model_maker/python/core/utils/quantization_test.py @@ -17,7 +17,6 @@ from mediapipe.model_maker.python.core.utils import quantization from mediapipe.model_maker.python.core.utils import test_util -from ai_edge_litert import interpreter as tfl_interpreter class QuantizationTest(tf.test.TestCase, parameterized.TestCase): @@ -60,7 +59,7 @@ def test_set_converter_with_quantization_from_int8_config(self): self.assertEqual(config.supported_ops, [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]) tflite_model = converter.convert() - interpreter = tfl_interpreter.Interpreter(model_content=tflite_model) + interpreter = tf.lite.Interpreter(model_content=tflite_model) self.assertEqual(interpreter.get_input_details()[0]['dtype'], tf.uint8) self.assertEqual(interpreter.get_output_details()[0]['dtype'], tf.uint8) @@ -83,7 +82,7 @@ def test_set_converter_with_quantization_from_float16_config(self): converter = config.set_converter_with_quantization(converter=converter) self.assertEqual(config.supported_types, [tf.float16]) tflite_model = converter.convert() - interpreter = tfl_interpreter.Interpreter(model_content=tflite_model) + interpreter = tf.lite.Interpreter(model_content=tflite_model) # The input and output are expected to be set to float32 by default. self.assertEqual(interpreter.get_input_details()[0]['dtype'], tf.float32) self.assertEqual(interpreter.get_output_details()[0]['dtype'], tf.float32) diff --git a/mediapipe/model_maker/requirements.txt b/mediapipe/model_maker/requirements.txt index 419d0166ca..3ce977b671 100644 --- a/mediapipe/model_maker/requirements.txt +++ b/mediapipe/model_maker/requirements.txt @@ -1,5 +1,4 @@ absl-py -ai-edge-litert-nightly mediapipe>=0.10.0 numpy<2 opencv-python diff --git a/mediapipe/model_maker/requirements_bazel.txt b/mediapipe/model_maker/requirements_bazel.txt index 003085ed07..fd6c421cf8 100644 --- a/mediapipe/model_maker/requirements_bazel.txt +++ b/mediapipe/model_maker/requirements_bazel.txt @@ -1,5 +1,4 @@ absl-py -ai-edge-litert-nightly numpy<2 opencv-python setuptools==70.3.0 # needed due to https://github.com/pypa/setuptools/issues/4487 diff --git a/mediapipe/model_maker/requirements_lock.txt b/mediapipe/model_maker/requirements_lock.txt index 62972c57da..83eb07a121 100644 --- a/mediapipe/model_maker/requirements_lock.txt +++ b/mediapipe/model_maker/requirements_lock.txt @@ -15,8 +15,6 @@ absl-py==1.4.0 # tensorflow-metadata # tensorflow-model-optimization # tf-slim -ai-edge-litert-nightly==1.0.1.dev20241022 - # via -r mediapipe/opensource_only/model_maker_requirements_bazel.txt array-record==0.5.1 # via tensorflow-datasets astunparse==1.6.3 @@ -50,9 +48,7 @@ etils[array-types,enp,epath,epy,etqdm,etree]==1.5.2 # array-record # tensorflow-datasets flatbuffers==24.3.25 - # via - # ai-edge-litert-nightly - # tensorflow + # via tensorflow fonttools==4.54.1 # via matplotlib fsspec==2024.9.0 @@ -124,7 +120,6 @@ ml-dtypes==0.3.2 numpy==1.26.4 # via # -r mediapipe/opensource_only/model_maker_requirements_bazel.txt - # ai-edge-litert-nightly # contourpy # etils # h5py