在今年的 I/O 大會上，我們很高興宣布推出一個基于嵌入向量的全新設備端搜索庫，讓您在幾毫秒的時間內即可從數百萬個數據樣本中快速找到相似的圖像、文本或音頻。

此搜索庫通過使用模型，將搜索查詢嵌入到表示查詢語義的高維向量中來執行搜索。隨后搜索庫使用 ScaNN（可擴容最近鄰算法）從預定義的數據庫中搜索相似項目。為將搜索庫應用到您的數據集，您需要使用 Model Maker Searcher API（教程）構建自定義 TFLite Searcher 模型，然后使用 Task Library Searcher API（視覺/文本）將其部署到設備上。

• ScaNN

  https://github.com/google-research/google-research/tree/master/scann

• 教程

  https://tensorflow.google.cn/lite/tutorials/model_maker_text_searcher

• 視覺

  https://tensorflow.google.cn/lite/inference_with_metadata/task_library/image_searcher

• 文本

  https://tensorflow.google.cn/lite/inference_with_metadata/task_library/text_searcher

例如，使用在 COCO 上訓練的 Searcher 模型，搜索查詢：“A passenger plane on the runway”，系統將返回以下圖像：

• COCO

  https://cocodataset.org/#home

圖 1：所有圖像均來自 COCO 2014 訓練和驗證數據集。圖像 1 由 Mark Jones Jr. 依據《版權歸屬許可證》提供。圖像 2 由 305 Seahill 依據《版權歸屬-禁止演繹許可證》提供。圖像 3 由 tataquax 依據《版權歸屬-相同方式共享許可證》提供。

在本文中，我們將向您介紹使用新 TensorFlow Lite Searcher Library 構建文本到圖像搜索功能的端到端示例（根據給定文本查詢檢索圖像）。以下是主要步驟：

1. 使用 COCO 數據集訓練用于圖像和文本查詢編碼的雙編碼器模型。

2. 使用 Model Maker Searcher API 創建文本到圖像 Searcher 模型。

3. 使用 Task Library Searcher API 檢索帶有文本查詢的圖像。

訓練雙編碼器模型

圖 2：用點積相似距離訓練雙編碼器模型。損失函數可為相關圖像和文本賦予更大的點積（陰影綠色方塊）

雙編碼器模型由圖像編碼器和文本編碼器組成。兩個編碼器分別將圖像和文本映射到高維空間中的嵌入向量。雙編碼器模型計算圖像和文本嵌入向量之間的點積，同時損失函數可為相關圖像和文本賦予更大的點積（更接近），而為不相關的圖像和文本賦予更小的點積（更遠）。

整個訓練過程受到了 CLIP 論文和本 Keras 示例的啟發。圖像編碼器是在預訓練 EfficientNet 模型的基礎上構建而成，而文本編碼器則是基于預訓練通用語句編碼器模型。

• CLIP

  https://arxiv.org/abs/2103.00020

• Keras 示例

  https://keras.io/examples/nlp/nl_image_search/

• EfficientNet

  https://hub.tensorflow.google.cn/google/imagenet/efficientnet_v2_imagenet21k_ft1k_s/feature_vector/2

• 通用語句編碼器

  https://hub.tensorflow.google.cn/google/universal-sentence-encoder-lite/2

系統隨后會將兩個編碼器的輸出投影到 128 維空間并進行 L2 歸一化。對于數據集，我們選擇使用 COCO，因為該數據集的訓練和驗證分塊會為每個圖像人工生成字幕。請查看配套的 Colab notebook，了解訓練過程的詳細信息。

• Colab notebook

  https://colab.sandbox.google.com/github/tensorflow/tflite-support/blob/master/tensorflow_lite_support/examples/colab/on_device_text_to_image_search_tflite.ipynb

雙編碼器模型可以從沒有字幕的數據庫中檢索圖像，因為在經過訓練后，圖像嵌入器可以直接從圖像中提取語義，而無需人工生成的字幕。

使用 Model Maker 創建文本

到圖像 Searcher 模型

圖 3：使用圖像編碼器生成圖像嵌入向量，并使用 Model Maker 創建 TFLite Searcher 模型

完成對雙編碼器模型的訓練后，我們可以使用它來創建 TFLite Searcher 模型，該模型可根據文本查詢，從圖像數據集中搜索最相關的圖像。模型創建分為以下三大步驟：

1. 使用 TensorFlow 圖像編碼器生成圖像數據集的嵌入向量。ScaNN 能夠搜索非常龐大的數據集，因此我們結合了 COCO 2014 的訓練和驗證分塊（總計超過 12.3 萬張圖像），以展示其搜索性能。相關代碼請查閱此處（1）。

2. 將 TensorFlow 文本編碼器模型轉換為 TFLite 格式。相關代碼請查閱此處（2）。

3. 使用 Model Maker，通過 TFLite 文本編碼器和使用以下代碼的圖像嵌入向量創建 TFLite Searcher 模型：

• 此處（1）

  https://colab.sandbox.google.com/github/tensorflow/tflite-support/blob/master/tensorflow_lite_support/examples/colab/on_device_text_to_image_search_tflite.ipynb#scrollTo=Bp0qBKkyu4jA

• 此處（2）

  https://colab.research.google.com/github/tensorflow/tflite-support/blob/master/tensorflow_lite_support/examples/colab/on_device_text_to_image_search_tflite.ipynb#scrollTo=6Dzye66Xc8vE

#Configure ScaNN options. See the API doc for how to configure ScaNN. 
scann_options = searcher.ScaNNOptions(
      distance_measure='dot_product',
      tree=searcher.Tree(num_leaves=351, num_leaves_to_search=4),
      score_ah=searcher.ScoreAH(1, anisotropic_quantization_threshold=0.2))

# Load the image embeddings and corresponding metadata if any.
data = searcher.DataLoader(tflite_embedder_path, image_embeddings, metadata)

# Create the TFLite Searcher model.
model = searcher.Searcher.create_from_data(data, scann_options)

# Export the TFLite Searcher model.
model.export(
      export_filename='searcher.tflite',
      userinfo='',
      export_format=searcher.ExportFormat.TFLITE)

請在此處查閱上方代碼中提到的 API doc。

• API doc

  https://tensorflow.google.cn/lite/api_docs/python/tflite_model_maker/searcher/ScaNNOptions

在創建 Searcher 模型時，Model Maker 利用 ScaNN 將嵌入向量編入索引。嵌入向量數據集首先被分為多個子集。在每個子集中，ScaNN 存儲嵌入向量的量化表征。在檢索時，ScaNN 會選擇一些最相關的分區，并按照快速近似距離對量化表征進行評分。這個過程既（通過量化）節省了模型大小又（通過分區選擇）實現了加速。請參閱深入研究資料，詳細了解 ScaNN 算法。

在上方示例中，我們將數據集劃分為 351 個分區（約是我們擁有的嵌入向量數量的平方根），并在檢索期間搜索其中的 4 個分區，即大約是數據集的 1%。我們還將 128 維浮點嵌入向量量化為 128 個 int8 值，以節省空間。

使用 Task Library 運行推理

圖 4：使用帶有 TFLite Searcher 模型的 Task Library 運行推理。推理接收查詢文本并返回最近鄰的元數據。我們可以在此找到對應的圖像

如要使用 Searcher 模型查詢圖像，您只需使用 Task Library 的幾行代碼即可，具體如下所示：

from tflite_support.task import text

# Initialize a TextSearcher object
searcher = text.TextSearcher.create_from_file('searcher.tflite')

# Search the input query
results = searcher.search(query_text)

# Show the results
for rank in range(len(results.nearest_neighbors)):
  print('Rank #', rank, ':')
  image_id = results.nearest_neighbors[rank].metadata
  print('image_id: ', image_id)
  print('distance: ', results.nearest_neighbors[rank].distance)
  show_image_by_id(image_id)

可以嘗試一下 Colab 的代碼。此外，歡迎查看更多信息，了解如何使用 Task Library Java 和 C++ API 集成模型，尤其是在 Android 上的用法。在 Pixel 6 上，每個查詢通常只需要 6 毫秒。

• 更多信息

  https://tensorflow.google.cn/lite/inference_with_metadata/task_library/text_searcher

以下是一些示例結果：

查詢：A man riding a bike

根據估算的相似距離對結果進行排序。以下是檢索到的圖像示例。請注意，我們僅會顯示附有圖像使用許可的圖像。

圖 5：所有圖像均來自 COCO 2014 訓練和驗證數據集。圖像 1 由 Reuel Mark Delez 依據《版權歸屬許可證》提供。圖像 2 由 Richard Masoner/Cyclelicious 依據《版權歸屬-相同方式共享許可證》提供。圖像 3 由 Julia 依據《版權歸屬-相同方式共享許可證》提供。圖像 4 由 Aaron Fulkerson 依據《版權歸屬-相同方式共享許可證》提供。圖像 5 由 Richard Masoner/Cyclelicious 依據《版權歸屬-相同方式共享許可證》提供。圖像 6 由 Richard Masoner/Cyclelicious 依據《版權歸屬-相同方式共享許可證》提供。

研究展望

我們將致力于啟用除圖像和文本之外的更多搜索類型，如音頻片段。

審核編輯 ：李倩