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

此搜索庫(kù)通過(guò)使用模型，將搜索查詢嵌入到表示查詢語(yǔ)義的高維向量中來(lái)執(zhí)行搜索。隨后搜索庫(kù)使用 ScaNN（可擴(kuò)容最近鄰算法）從預(yù)定義的數(shù)據(jù)庫(kù)中搜索相似項(xiàng)目。為將搜索庫(kù)應(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：所有圖像均來(lái)自 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)積（陰影綠色方塊）

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

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

• 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

• 通用語(yǔ)句編碼器

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

系統(tǒng)隨后會(huì)將兩個(gè)編碼器的輸出投影到 128 維空間并進(jìn)行 L2 歸一化。對(duì)于數(shù)據(jù)集，我們選擇使用 COCO，因?yàn)樵摂?shù)據(jù)集的訓(xùn)練和驗(yàn)證分塊會(huì)為每個(gè)圖像人工生成字幕。請(qǐng)查看配套的 Colab notebook，了解訓(xùn)練過(guò)程的詳細(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ù)庫(kù)中檢索圖像，因?yàn)樵诮?jīng)過(guò)訓(xùn)練后，圖像嵌入器可以直接從圖像中提取語(yǔ)義，而無(wú)需人工生成的字幕。

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

到圖像 Searcher 模型

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

完成對(duì)雙編碼器模型的訓(xùn)練后，我們可以使用它來(lái)創(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)證分塊（總計(jì)超過(guò) 12.3 萬(wàn)張圖像），以展示其搜索性能。相關(guān)代碼請(qǐng)查閱此處（1）。

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

3. 使用 Model Maker，通過(guò) 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)

請(qǐng)?jiān)诖颂幉殚喩戏酱a中提到的 API doc。

• API doc

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

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

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

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

圖 4：使用帶有 TFLite Searcher 模型的 Task Library 運(yùn)行推理。推理接收查詢文本并返回最近鄰的元數(shù)據(jù)。我們可以在此找到對(duì)應(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 上，每個(gè)查詢通常只需要 6 毫秒。

• 更多信息

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

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

查詢：A man riding a bike

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

圖 5：所有圖像均來(lái)自 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)歸屬-相同方式共享許可證》提供。

研究展望

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

審核編輯 ：李倩