LiteRT 使用入门

本指南介绍了在设备上运行 LiteRT(Lite Runtime 的简称)模型以根据输入数据进行预测的过程。这是通过 LiteRT 解释器实现的,该解释器使用静态图排序和自定义(不太动态)内存分配器来确保最小的加载、初始化和执行延迟时间。

LiteRT 推理通常遵循以下步骤:

  1. 加载模型:将 .tflite 模型加载到内存中,其中包含模型的执行图。

  2. 转换数据:将输入数据转换为预期格式和维度。模型的原始输入数据通常与模型所需的输入数据格式不匹配。例如,您可能需要调整图片大小或更改图片格式,以使其与模型兼容。

  3. 运行推理:执行 LiteRT 模型以进行预测。此步骤涉及使用 LiteRT API 执行模型。它涉及几个步骤,例如构建解释器和分配张量。

  4. 解读输出:以有意义的方式解读输出张量,以便在应用中加以利用。例如,模型可能仅返回概率列表。您可以自行将概率映射到相关类别并设置输出格式。

本指南介绍了如何访问 LiteRT 解释器,以及如何使用 C++、Java 和 Python 执行推理。

支持的平台

TensorFlow 推断 API 适用于大多数常见的移动和嵌入式平台,例如 AndroidiOSLinux,并提供多种编程语言版本

在大多数情况下,API 设计都体现了对性能的偏好,而不是易用性。LiteRT 旨在在小型设备上实现快速推理,因此 API 会避免不必要的复制,但会牺牲便利性。

在所有库中,LiteRT API 都允许您加载模型、提供输入和检索推理输出。

Android 平台

在 Android 上,可以使用 Java 或 C++ API 执行 LiteRT 推理。Java API 非常方便,可以直接在 Android Activity 类中使用。C++ API 更加灵活且速度更快,但可能需要编写 JNI 封装容器才能在 Java 层和 C++ 层之间移动数据。

如需了解详情,请参阅 C++Java 部分,或按照 Android 快速入门操作。

iOS 平台

在 iOS 上,LiteRT 可在 SwiftObjective-C iOS 库中使用。您还可以直接在 Objective-C 代码中使用 C API

请参阅 SwiftObjective-CC API 部分,或按照 iOS 快速入门操作。

Linux 平台

在 Linux 平台上,您可以使用 C++ 中提供的 LiteRT API 运行推理。

加载并运行模型

加载和运行 LiteRT 模型涉及以下步骤:

  1. 将模型加载到内存中。
  2. 基于现有模型构建 Interpreter
  3. 设置输入张量值。
  4. 调用推理。
  5. 输出张量值。

Android (Java)

用于通过 LiteRT 运行推理的 Java API 主要设计用于 Android,因此可作为 Android 库依赖项使用:com.google.ai.edge.litert

在 Java 中,您将使用 Interpreter 类来加载模型并驱动模型推理。在许多情况下,这可能是您唯一需要的 API。

您可以使用 FlatBuffers (.tflite) 文件初始化 Interpreter

public Interpreter(@NotNull File modelFile);

或者使用 MappedByteBuffer

public Interpreter(@NotNull MappedByteBuffer mappedByteBuffer);

在这两种情况下,您都必须提供有效的 LiteRT 模型,否则 API 会抛出 IllegalArgumentException。如果您使用 MappedByteBuffer 初始化 Interpreter,则在整个 Interpreter 生命周期内,该值必须保持不变。

对模型运行推理的首选方式是使用签名 - 适用于从 TensorFlow 2.5 开始转换的模型

try (Interpreter interpreter = new Interpreter(file_of_tensorflowlite_model)) {
  Map<String, Object> inputs = new HashMap<>();
  inputs.put("input_1", input1);
  inputs.put("input_2", input2);
  Map<String, Object> outputs = new HashMap<>();
  outputs.put("output_1", output1);
  interpreter.runSignature(inputs, outputs, "mySignature");
}

runSignature 方法采用三个实参:

  • 输入:从签名中的输入名称到输入对象的输入映射。

  • 输出:用于从签名中的输出名称到输出数据的输出映射的映射。

  • 签名名称(可选):签名名称(如果模型具有单个签名,则可以留空)。

在模型没有已定义的签名时运行推理的另一种方式。 只需调用 Interpreter.run() 即可。例如:

try (Interpreter interpreter = new Interpreter(file_of_a_tensorflowlite_model)) {
  interpreter.run(input, output);
}

run() 方法仅接受一个输入并仅返回一个输出。因此,如果您的模型有多个输入或多个输出,请改用:

interpreter.runForMultipleInputsOutputs(inputs, map_of_indices_to_outputs);

在这种情况下,inputs 中的每个条目都对应一个输入张量,而 map_of_indices_to_outputs 将输出张量的索引映射到相应的输出数据。

在这两种情况下,张量指数都应与您在创建模型时为 LiteRT 转换器提供的值相对应。请注意,input 中张量的顺序必须与传递给 LiteRT 转换器的顺序一致。

Interpreter 类还提供了一些便捷的函数,让您可以使用操作名称获取任何模型输入或输出的索引:

public int getInputIndex(String opName);
public int getOutputIndex(String opName);

如果 opName 不是模型中的有效操作,则会抛出 IllegalArgumentException

另请注意,Interpreter 拥有资源。为避免内存泄漏,必须在使用后通过以下方式释放资源:

interpreter.close();

如需查看使用 Java 的示例项目,请参阅 Android 对象检测示例应用

支持的数据类型

如需使用 LiteRT,输入和输出张量的数据类型必须是以下原始类型之一:

  • float
  • int
  • long
  • byte

系统还支持 String 类型,但它们的编码方式与原始类型不同。具体而言,字符串张量的形状决定了张量中字符串的数量和排列方式,而每个元素本身都是一个可变长度的字符串。从这个意义上讲,Tensor 的(字节)大小无法仅根据形状和类型来计算,因此字符串无法作为单个扁平的 ByteBuffer 实参提供。

如果使用其他数据类型(包括 IntegerFloat 等装箱类型),则会抛出 IllegalArgumentException

输入

每个输入都应该是受支持的原始类型的数组或多维数组,或者大小合适的原始 ByteBuffer。如果输入是数组或多维数组,则关联的输入张量将在推理时隐式调整为数组的维度。如果输入是 ByteBuffer,调用方应先手动调整关联的输入张量的大小(通过 Interpreter.resizeInput()),然后再运行推理。

使用 ByteBuffer 时,最好使用直接字节缓冲区,因为这样可让 Interpreter 避免不必要的复制。如果 ByteBuffer 是直接字节缓冲区,则其顺序必须为 ByteOrder.nativeOrder()。在用于模型推理后,它必须保持不变,直到模型推理完成。

输出

每个输出都应该是受支持的原始类型的数组或多维数组,或者大小合适的 ByteBuffer。请注意,某些模型具有动态输出,其中输出张量的形状可能会因输入而异。使用现有的 Java 推断 API 无法直接处理此问题,但计划的扩展功能将使此问题得以解决。

iOS (Swift)

Swift API 可从 Cocoapods 的 TensorFlowLiteSwift Pod 中获取。

首先,您需要导入 TensorFlowLite 模块。

import TensorFlowLite
// Getting model path
guard
  let modelPath = Bundle.main.path(forResource: "model", ofType: "tflite")
else {
  // Error handling...
}

do {
  // Initialize an interpreter with the model.
  let interpreter = try Interpreter(modelPath: modelPath)

  // Allocate memory for the model's input `Tensor`s.
  try interpreter.allocateTensors()

  let inputData: Data  // Should be initialized

  // input data preparation...

  // Copy the input data to the input `Tensor`.
  try self.interpreter.copy(inputData, toInputAt: 0)

  // Run inference by invoking the `Interpreter`.
  try self.interpreter.invoke()

  // Get the output `Tensor`
  let outputTensor = try self.interpreter.output(at: 0)

  // Copy output to `Data` to process the inference results.
  let outputSize = outputTensor.shape.dimensions.reduce(1, {x, y in x * y})
  let outputData =
        UnsafeMutableBufferPointer<Float32>.allocate(capacity: outputSize)
  outputTensor.data.copyBytes(to: outputData)

  if (error != nil) { /* Error handling... */ }
} catch error {
  // Error handling...
}

iOS (Objective-C)

Objective-C API 可在 Cocoapods 的 LiteRTObjC Pod 中使用。

首先,您需要导入 TensorFlowLiteObjC 模块。

@import TensorFlowLite;
NSString *modelPath = [[NSBundle mainBundle] pathForResource:@"model"
                                                      ofType:@"tflite"];
NSError *error;

// Initialize an interpreter with the model.
TFLInterpreter *interpreter = [[TFLInterpreter alloc] initWithModelPath:modelPath
                                                                  error:&error];
if (error != nil) { /* Error handling... */ }

// Allocate memory for the model's input `TFLTensor`s.
[interpreter allocateTensorsWithError:&error];
if (error != nil) { /* Error handling... */ }

NSMutableData *inputData;  // Should be initialized
// input data preparation...

// Get the input `TFLTensor`
TFLTensor *inputTensor = [interpreter inputTensorAtIndex:0 error:&error];
if (error != nil) { /* Error handling... */ }

// Copy the input data to the input `TFLTensor`.
[inputTensor copyData:inputData error:&error];
if (error != nil) { /* Error handling... */ }

// Run inference by invoking the `TFLInterpreter`.
[interpreter invokeWithError:&error];
if (error != nil) { /* Error handling... */ }

// Get the output `TFLTensor`
TFLTensor *outputTensor = [interpreter outputTensorAtIndex:0 error:&error];
if (error != nil) { /* Error handling... */ }

// Copy output to `NSData` to process the inference results.
NSData *outputData = [outputTensor dataWithError:&error];
if (error != nil) { /* Error handling... */ }

Objective-C 代码中的 C API

Objective-C API 不支持委托。为了在 Objective-C 代码中使用委托,您需要直接调用底层 C API

#include "tensorflow/lite/c/c_api.h"
TfLiteModel* model = TfLiteModelCreateFromFile([modelPath UTF8String]);
TfLiteInterpreterOptions* options = TfLiteInterpreterOptionsCreate();

// Create the interpreter.
TfLiteInterpreter* interpreter = TfLiteInterpreterCreate(model, options);

// Allocate tensors and populate the input tensor data.
TfLiteInterpreterAllocateTensors(interpreter);
TfLiteTensor* input_tensor =
    TfLiteInterpreterGetInputTensor(interpreter, 0);
TfLiteTensorCopyFromBuffer(input_tensor, input.data(),
                           input.size() * sizeof(float));

// Execute inference.
TfLiteInterpreterInvoke(interpreter);

// Extract the output tensor data.
const TfLiteTensor* output_tensor =
    TfLiteInterpreterGetOutputTensor(interpreter, 0);
TfLiteTensorCopyToBuffer(output_tensor, output.data(),
                         output.size() * sizeof(float));

// Dispose of the model and interpreter objects.
TfLiteInterpreterDelete(interpreter);
TfLiteInterpreterOptionsDelete(options);
TfLiteModelDelete(model);

C++

用于通过 LiteRT 运行推理的 C++ API 与 Android、iOS 和 Linux 平台兼容。iOS 上的 C++ API 仅在使用 bazel 时可用。

在 C++ 中,模型存储在 FlatBufferModel 类中。它封装了 LiteRT 模型,您可以根据模型的存储位置,通过几种不同的方式来构建它:

class FlatBufferModel {
  // Build a model based on a file. Return a nullptr in case of failure.
  static std::unique_ptr<FlatBufferModel> BuildFromFile(
      const char* filename,
      ErrorReporter* error_reporter);

  // Build a model based on a pre-loaded flatbuffer. The caller retains
  // ownership of the buffer and should keep it alive until the returned object
  // is destroyed. Return a nullptr in case of failure.
  static std::unique_ptr<FlatBufferModel> BuildFromBuffer(
      const char* buffer,
      size_t buffer_size,
      ErrorReporter* error_reporter);
};

现在,您已将模型作为 FlatBufferModel 对象,可以使用 Interpreter 执行该模型。多个 Interpreter 可以同时使用一个 FlatBufferModel

以下代码段显示了 Interpreter API 的重要部分。请注意:

  • 张量由整数表示,以避免字符串比较(以及对字符串库的任何固定依赖关系)。
  • 不得从并发线程访问解释器。
  • 必须在调整张量大小后立即调用 AllocateTensors(),以触发输入和输出张量的内存分配。

使用 C++ 的 LiteRT 最简单的用法如下所示:

// Load the model
std::unique_ptr<tflite::FlatBufferModel> model =
    tflite::FlatBufferModel::BuildFromFile(filename);

// Build the interpreter
tflite::ops::builtin::BuiltinOpResolver resolver;
std::unique_ptr<tflite::Interpreter> interpreter;
tflite::InterpreterBuilder(*model, resolver)(&interpreter);

// Resize input tensors, if needed.
interpreter->AllocateTensors();

float* input = interpreter->typed_input_tensor<float>(0);
// Fill `input`.

interpreter->Invoke();

float* output = interpreter->typed_output_tensor<float>(0);

如需查看更多示例代码,请参阅 minimal.cclabel_image.cc

Python

用于运行推理的 Python API 使用 Interpreter 加载模型并运行推理。

安装 LiteRT 软件包:

$ python3 -m pip install ai-edge-litert

导入 LiteRT 解释器

from ai_edge_litert.interpreter import Interpreter
Interpreter = Interpreter(model_path=args.model.file)

以下示例展示了如何使用 Python 解释器加载 FlatBuffers (.tflite) 文件并使用随机输入数据运行推理:

如果您要从具有已定义 SignatureDef 的 SavedModel 进行转换,建议使用此示例。

class TestModel(tf.Module):
  def __init__(self):
    super(TestModel, self).__init__()

  @tf.function(input_signature=[tf.TensorSpec(shape=[1, 10], dtype=tf.float32)])
  def add(self, x):
    '''
    Simple method that accepts single input 'x' and returns 'x' + 4.
    '''
    # Name the output 'result' for convenience.
    return {'result' : x + 4}

SAVED_MODEL_PATH = 'content/saved_models/test_variable'
TFLITE_FILE_PATH = 'content/test_variable.tflite'

# Save the model
module = TestModel()
# You can omit the signatures argument and a default signature name will be
# created with name 'serving_default'.
tf.saved_model.save(
    module, SAVED_MODEL_PATH,
    signatures={'my_signature':module.add.get_concrete_function()})

# Convert the model using TFLiteConverter
converter = tf.lite.TFLiteConverter.from_saved_model(SAVED_MODEL_PATH)
tflite_model = converter.convert()
with open(TFLITE_FILE_PATH, 'wb') as f:
  f.write(tflite_model)

# Load the LiteRT model in LiteRT Interpreter
from ai_edge_litert.interpreter import Interpreter
interpreter = Interpreter(TFLITE_FILE_PATH)

# There is only 1 signature defined in the model,
# so it will return it by default.
# If there are multiple signatures then we can pass the name.
my_signature = interpreter.get_signature_runner()

# my_signature is callable with input as arguments.
output = my_signature(x=tf.constant([1.0], shape=(1,10), dtype=tf.float32))
# 'output' is dictionary with all outputs from the inference.
# In this case we have single output 'result'.
print(output['result'])

如果模型未定义 SignatureDefs,则为另一个示例。

import numpy as np
import tensorflow as tf

# Load the LiteRT model and allocate tensors.
from ai_edge_litert.interpreter import Interpreter
interpreter = Interpreter(TFLITE_FILE_PATH)
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Test the model on random input data.
input_shape = input_details[0]['shape']
input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)

interpreter.invoke()

# The function `get_tensor()` returns a copy of the tensor data.
# Use `tensor()` in order to get a pointer to the tensor.
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

除了将模型作为预转换的 .tflite 文件加载之外,您还可以将代码与 LiteRT 编译器结合使用,从而将 Keras 模型转换为 LiteRT 格式,然后运行推理:

import numpy as np
import tensorflow as tf

img = tf.keras.Input(shape=(64, 64, 3), name="img")
const = tf.constant([1., 2., 3.]) + tf.constant([1., 4., 4.])
val = img + const
out = tf.identity(val, name="out")

# Convert to LiteRT format
converter = tf.lite.TFLiteConverter.from_keras_model(tf.keras.models.Model(inputs=[img], outputs=[out]))
tflite_model = converter.convert()

# Load the LiteRT model and allocate tensors.
from ai_edge_litert.interpreter import Interpreter
interpreter = Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

# Continue to get tensors and so forth, as shown above...

如需查看更多 Python 示例代码,请参阅 label_image.py

使用动态形状模型运行推理

如果您想运行具有动态输入形状的模型,请在运行推理之前调整输入形状的大小。否则,Tensorflow 模型中的 None 形状将被 LiteRT 模型中的 1 占位符替换。

以下示例展示了如何在不同语言中调整输入形状大小,然后再运行推理。所有示例都假设输入形状定义为 [1/None, 10],并且需要调整为 [3, 10]

C++ 示例:

// Resize input tensors before allocate tensors
interpreter->ResizeInputTensor(/*tensor_index=*/0, std::vector<int>{3,10});
interpreter->AllocateTensors();

Python 示例:

# Load the LiteRT model in LiteRT Interpreter
from ai_edge_litert.interpreter import Interpreter
interpreter = Interpreter(model_path=TFLITE_FILE_PATH)

# Resize input shape for dynamic shape model and allocate tensor
interpreter.resize_tensor_input(interpreter.get_input_details()[0]['index'], [3, 10])
interpreter.allocate_tensors()

# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()