python/tflite_micro/runtime_test.py

# Copyright 2022 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.
# ==============================================================================
"""Basic Python test for the TFLM interpreter"""

# Steps to debug with gdb:
# 1. bazel build python/tflite_micro:runtime_test
# 2. gdb python
# 3. (gdb) run bazel-out/k8-fastbuild/bin/python/tflite_micro/runtime_test

import gc
import weakref
import numpy as np
import tensorflow as tf

from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
from tflite_micro.tensorflow.lite.micro.testing import generate_test_models
from tflite_micro.python.tflite_micro import runtime


class ConvModelTests(test_util.TensorFlowTestCase):
  filename = "/tmp/interpreter_test_conv_model.tflite"
  input_shape = (1, 16, 16, 1)
  output_shape = (1, 10)

  def testInitErrorHandling(self):
    with self.assertRaisesWithPredicateMatch(ValueError,
                                             "Invalid model file path"):
      runtime.Interpreter.from_file("wrong.tflite")

  def testInput(self):
    model_data = generate_test_models.generate_conv_model(False)
    tflm_interpreter = runtime.Interpreter.from_bytes(model_data)

    data_x = np.random.randint(-127, 127, self.input_shape, dtype=np.int8)
    tflm_interpreter.set_input(data_x, 0)

    # Test input tensor details
    input_details = tflm_interpreter.get_input_details(0)
    self.assertAllEqual(input_details["shape"], self.input_shape)
    # Single channel int8 quantization
    self.assertEqual(input_details["dtype"], np.int8)
    self.assertEqual(len(input_details["quantization_parameters"]["scales"]),
                     1)
    self.assertEqual(
        input_details["quantization_parameters"]["quantized_dimension"], 0)
    # TODO(b/247808903): check only the types here to make sure that all arrays are properly set up.
    self.assertEqual(input_details["quantization_parameters"]["scales"].dtype,
                     np.float32)
    self.assertEqual(
        input_details["quantization_parameters"]["zero_points"].dtype,
        np.int32)

  def testInputErrorHandling(self):
    model_data = generate_test_models.generate_conv_model(True, self.filename)
    tflm_interpreter = runtime.Interpreter.from_bytes(model_data)

    data_x = np.random.randint(-127, 127, self.input_shape, dtype=np.int8)
    # Try to access out of bound data
    with self.assertRaisesWithPredicateMatch(IndexError,
                                             "Tensor is out of bound"):
      tflm_interpreter.set_input(data_x, 1)
    # Pass data with wrong dimension
    with self.assertRaisesWithPredicateMatch(ValueError,
                                             "Dimension mismatch."):
      reshaped_data = data_x.reshape((1, 16, 16, 1, 1))
      tflm_interpreter.set_input(reshaped_data, 0)
    # Pass data with wrong dimension in one axis
    with self.assertRaisesWithPredicateMatch(ValueError,
                                             "Dimension mismatch."):
      reshaped_data = data_x.reshape((1, 2, 128, 1))
      tflm_interpreter.set_input(reshaped_data, 0)
    # Pass data with wrong type
    with self.assertRaisesWithPredicateMatch(ValueError, "Got value of type"):
      float_data = data_x.astype(np.float32)
      tflm_interpreter.set_input(float_data, 0)
    # Reach wrong details
    with self.assertRaisesWithPredicateMatch(IndexError,
                                             "Tensor is out of bound"):
      tflm_interpreter.get_input_details(1)

  def testOutput(self):
    model_data = generate_test_models.generate_conv_model(True, self.filename)
    tflm_interpreter = runtime.Interpreter.from_bytes(model_data)

    # Initial output values are all 0
    output = tflm_interpreter.get_output(0)
    init_output = np.zeros(self.output_shape)
    self.assertAllEqual(output, init_output)

    # Test the output tensor details
    output_details = tflm_interpreter.get_output_details(0)
    self.assertAllEqual(output_details["shape"], self.output_shape)
    # Single channel int8 quantization
    self.assertEqual(output_details["dtype"], np.int8)
    self.assertEqual(len(output_details["quantization_parameters"]["scales"]),
                     1)
    self.assertEqual(
        output_details["quantization_parameters"]["quantized_dimension"], 0)
    # TODO(b/247808903): check only the types here to make sure that all arrays are properly set up.
    self.assertEqual(output_details["quantization_parameters"]["scales"].dtype,
                     np.float32)
    self.assertEqual(
        output_details["quantization_parameters"]["zero_points"].dtype,
        np.int32)

  def testOutputErrorHandling(self):
    model_data = generate_test_models.generate_conv_model(True, self.filename)
    tflm_interpreter = runtime.Interpreter.from_bytes(model_data)
    # Try to access out of bound data
    with self.assertRaisesWithPredicateMatch(IndexError,
                                             "Tensor is out of bound"):
      tflm_interpreter.get_output(1)
    with self.assertRaisesWithPredicateMatch(IndexError,
                                             "Tensor is out of bound"):
      tflm_interpreter.get_output_details(1)

  def testCompareWithTFLite(self):
    model_data = generate_test_models.generate_conv_model(True, self.filename)

    # TFLM interpreter
    tflm_interpreter = runtime.Interpreter.from_bytes(model_data)

    # TFLite interpreter
    tflite_interpreter = tf.lite.Interpreter(
        model_content=model_data,
        experimental_op_resolver_type=\
        tf.lite.experimental.OpResolverType.BUILTIN_REF)
    tflite_interpreter.allocate_tensors()
    tflite_output_details = tflite_interpreter.get_output_details()[0]
    tflite_input_details = tflite_interpreter.get_input_details()[0]

    num_steps = 100
    for i in range(0, num_steps):
      # Create random input
      data_x = np.random.randint(-127, 127, self.input_shape, dtype=np.int8)

      # Run inference on TFLite
      tflite_interpreter.set_tensor(tflite_input_details["index"], data_x)
      tflite_interpreter.invoke()
      tflite_output = tflite_interpreter.get_tensor(
          tflite_output_details["index"])

      # Run inference on TFLM
      tflm_interpreter.set_input(data_x, 0)
      tflm_interpreter.invoke()
      tflm_output = tflm_interpreter.get_output(0)

      # Check that TFLM output has correct metadata
      self.assertDTypeEqual(tflm_output, np.int8)
      self.assertEqual(tflm_output.shape, self.output_shape)
      self.assertAllEqual(tflite_output, tflm_output)

  def _helperModelFromFileAndBufferEqual(self):
    model_data = generate_test_models.generate_conv_model(True, self.filename)

    file_interpreter = runtime.Interpreter.from_file(self.filename)
    bytes_interpreter = runtime.Interpreter.from_bytes(model_data)

    num_steps = 100
    for i in range(0, num_steps):
      data_x = np.random.randint(-127, 127, self.input_shape, dtype=np.int8)

      file_interpreter.set_input(data_x, 0)
      file_interpreter.invoke()
      file_output = file_interpreter.get_output(0)

      bytes_interpreter.set_input(data_x, 0)
      bytes_interpreter.invoke()
      bytes_output = bytes_interpreter.get_output(0)

      self.assertDTypeEqual(file_output, np.int8)
      self.assertEqual(file_output.shape, self.output_shape)
      self.assertDTypeEqual(bytes_output, np.int8)
      self.assertEqual(bytes_output.shape, self.output_shape)
      # Same interpreter and model, should expect all equal
      self.assertAllEqual(file_output, bytes_output)

  def testModelFromFileAndBufferEqual(self):
    self._helperModelFromFileAndBufferEqual()

  def testMultipleInterpreters(self):
    model_data = generate_test_models.generate_conv_model(False)

    interpreters = [
        runtime.Interpreter.from_bytes(model_data) for i in range(10)
    ]

    num_steps = 100
    for i in range(0, num_steps):
      data_x = np.random.randint(-127, 127, self.input_shape, dtype=np.int8)

      prev_output = None
      for interpreter in interpreters:
        interpreter.set_input(data_x, 0)
        interpreter.invoke()
        output = interpreter.get_output(0)
        if prev_output is None:
          prev_output = output

        self.assertDTypeEqual(output, np.int8)
        self.assertEqual(output.shape, self.output_shape)
        self.assertAllEqual(output, prev_output)

  def _helperNoop(self):
    pass

  def _helperOutputTensorMemoryLeak(self):
    interpreter = runtime.Interpreter.from_file(self.filename)
    int_ref = weakref.finalize(interpreter, self._helperNoop)
    some_output = interpreter.get_output(0)
    output_ref = weakref.finalize(some_output, self._helperNoop)
    return (int_ref, output_ref)

  def testOutputTensorMemoryLeak(self):
    generate_test_models.generate_conv_model(True, self.filename)

    int_ref, output_ref = self._helperOutputTensorMemoryLeak()
    # Output obtained in the helper function should be out of scope now, perform
    # garbage collection and check that the weakref is dead. If it's still
    # alive, it means that the output's reference count isn't 0 by garbage
    # collection. Since it's already out of scope, this means a memory leak.
    #
    # An example of how this could be true is if there's an additional
    # reference increment (e.g. `Py_INCREF` or `py::cast`` instead of
    # `py::reinterpret_steal``) somewhere in the C++ code.
    gc.collect()
    self.assertFalse(int_ref.alive)
    self.assertFalse(output_ref.alive)

  # TODO(b/240162715): Add a test case to register a custom OP

  def testMalformedCustomOps(self):
    model_data = generate_test_models.generate_conv_model(False)
    custom_op_registerers = [("wrong", "format")]
    with self.assertRaisesWithPredicateMatch(ValueError,
                                             "must be a list of strings"):
      interpreter = runtime.Interpreter.from_bytes(model_data,
                                                   custom_op_registerers)

    custom_op_registerers = "WrongFormat"
    with self.assertRaisesWithPredicateMatch(ValueError,
                                             "must be a list of strings"):
      interpreter = runtime.Interpreter.from_bytes(model_data,
                                                   custom_op_registerers)

  def testNonExistentCustomOps(self):
    model_data = generate_test_models.generate_conv_model(False)
    custom_op_registerers = ["SomeRandomOp"]
    with self.assertRaisesWithPredicateMatch(
        RuntimeError, "TFLM could not register custom op via SomeRandomOp"):
      interpreter = runtime.Interpreter.from_bytes(model_data,
                                                   custom_op_registerers)


if __name__ == "__main__":
  test.main()