On-device Text-to-Image Search with TensorFlow Lite Searcher Library

(Scalable Nearest Neighbors) to search for similar items from a predefined database. In order to apply it to your dataset, you need to use Model Maker Searcher API () to build a custom TFLite Searcher model, and then deploy it onto devices using Task Library Searcher API ().

For example, with the Searcher model trained on Figure 1: All images are from COCO 2014 train and validation datasets. . . .

In this post, we will walk you through an end-to-end example of building a text-to-image search feature (retrieve the images given textual queries) using the new TensorFlow Lite Searcher Library. Here are the major steps:

1. Train a dual encoder model for image and text query encoding using the Figure 2: Train the dual encoder model with dot product similarity distance. The loss encourages related images and text to have larger dot products (the shaded green squares).

   The dual encoder model consists of an image encoder and a text encoder. The two encoders map the images and text, respectively, to embeddings in a high-dimensional space. The model computes the dot product between the image and text embeddings, and the loss encourages relevant image and text to have larger dot product (closer), and unrelated ones to have smaller dot product (farther apart).

   The training procedure is inspired by the . The image encoder is based on a pre-trained model. The outputs from both encoders are then projected to a 128 dimensional space and are L2 normalized. For the dataset, we chose to use for the details of the training process.

   The dual encoder model makes it possible to retrieve images from a database without captions because once trained, the image embedder can directly extract the semantic meaning from the image without any need for human-generated captions.

   Create the text-to-image Searcher model using Model Maker

   . Convert the TensorFlow text encoder model into TFLite format. See the code to create the TFLite Searcher model from the TFLite text encoder and the image embeddings using the code below: PYTHON Kopieren #Configure ScaNN options. See the When creating the Searcher model, Model Maker leverages ScaNN to index the embedding vectors. The embedding dataset is first partitioned into multiple subsets. In each of the subsets, ScaNN stores the quantized representation of the embedding vectors. At retrieval time, ScaNN selects a few most relevant partitions and scores the quantized representations with fast, approximate distances. This procedure saves both the model size (through quantization) and achieves speed up (through partition selection). See the

   Figure 4: Run inference using Task Library with the TFLite Searcher model. It takes the query text and returns the top neighbor’s metadata. From there we can find the corresponding images.

   To query images using the Searcher model, you only need a couple of lines of code like the following using . Also, see Figure 5: All images are from COCO 2014 train and validation datasets. . . . . . .

   Future work

   We’ll be working on enabling more search types beyond image and text, such as audio clips.

   Contact [email protected] if you want to leave any feedback. Our goal is to make on-device ML even easier for you and we value your input!

   Acknowledgements

   We would like to thank Khanh LeViet, Chuo-Ling Chang, Ruiqi Guo, Lawrence Chan, Laurence Moroney, Yu-Cheng Ling, Matthias Grundmann, as well as Robby Neale, Chung-Ching Chang‎, Tom Small and Khalid Salama for their active support of this work. We would also like to thank the entire ScaNN team: David Simcha, Erik Lindgren, Felix Chern, Phil Sun and Sanjiv Kumar.

   Vollständiger Original-Bericht

   Ausführliche Details, Code-Beispiele & Hersteller-Stellungnahme auf blog.tensorflow.org.

   ↗ Original-Artikel auf blog.tensorflow.org lesen