基于嵌入向量的全新設(shè)備端搜索庫

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

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

• 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 上訓(xùn)練的 Searcher 模型，搜索查詢：“A passenger plane on the runway”，系統(tǒng)將返回以下圖像：

• COCO

  https://cocodataset.org/#home

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

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

1. 使用 COCO 數(shù)據(jù)集訓(xùn)練用于圖像和文本查詢編碼的雙編碼器模型。

2. 使用 Model Maker Searcher API 創(chuàng)建文本到圖像 Searcher 模型。

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

訓(xùn)練雙編碼器模型

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

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

整個訓(xùn)練過程受到了 CLIP 論文和本 Keras 示例的啟發(fā)。圖像編碼器是在預(yù)訓(xùn)練 EfficientNet 模型的基礎(chǔ)上構(gòu)建而成，而文本編碼器則是基于預(yù)訓(xùn)練通用語句編碼器模型。

• 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

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

• 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

雙編碼器模型可以從沒有字幕的數(shù)據(jù)庫中檢索圖像，因?yàn)樵诮?jīng)過訓(xùn)練后，圖像嵌入器可以直接從圖像中提取語義，而無需人工生成的字幕。

使用 Model Maker 創(chuàng)建文本

到圖像 Searcher 模型

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

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

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

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

3. 使用 Model Maker，通過 TFLite 文本編碼器和使用以下代碼的圖像嵌入向量創(chuàng)建 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

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

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

使用 Task Library 運(yùn)行推理

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

如要使用 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

以下是一些示例結(jié)果：

查詢：A man riding a bike

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

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

研究展望

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

審核編輯 ：李倩