Как работает поиск изображений в Dropbox

Если вам нужно найти фотографию, сделанную на пикнике несколько лет назад, вряд ли вы помните имя, которое камера автоматически присвоила файлу в момент съёмки, например, 2017-07-04 12.37.54.jpg. Вы просматриваете всё подряд — фотографии, их эскизы, пытаетесь определить объекты или признаки искомого — и не важно, ищете ли вы потерянное фото или хотите подыскать в архивах приличный снимок для презентации нового проекта.

Вот было бы здорово, если бы Dropbox мог самостоятельно просматривать все изображения и выбирать их них те, которые лучше всего соответствуют заданным в описании словам! Именно эту задачу мы поставили перед собой, создавая функцию поиска изображений.

К старту курса о машинном и глубоком обучении мы решили поделиться переводом о том, как текстовый поиск по изображениям в Dropbox работает изнутри, на каких наборах данных обучалась решающая задачу модель, как комбинировались методы, а также о том, какие Dropbox получила результаты и над чем работает сегодня.

В этой статье мы расскажем о созданном нами методе поиска содержимого в изображениях, основанном на технологии машинного обучения, и о том, как мы успешно интегрировали наш метод в существующую инфраструктуру поиска в Dropbox.

Наш метод

Формулируем задачу поиска изображений: необходимо найти функцию релевантности, которая принимала бы на входе (текстовый) запрос q и изображение j, а на выходе выдавала бы оценку релевантности s в баллах, показывающую, насколько точно изображение соответствует запросу:

s = f(q, j).

Если пользователю нужно найти какое-либо изображение, мы применяем эту функцию ко всем изображениям, и функция выдаёт отсортированный по степени приближения список изображений с балльной оценкой выше порогового значения. Мы создавали такую функцию на базе двух ключевых разработок в области машинного обучения: точной классификации изображений и векторов слов.

Классификация изображений

Классификатор изображений считывает изображение и выдаёт балльный список категорий, описывающих его содержание. Более высокие баллы указывают на большую вероятность того, что изображение относится к данной категории. 

Категории могут быть следующими:

• конкретные объекты на изображении, например, дерево или человек;

• общие дескрипторы сцены, например, на улице или свадьба;

• характеристики самого изображения, например, чёрно-белое или крупный план.

За последнее десятилетие был сделан огромный шаг вперёд в области в классификации изображений с помощью свёрточных нейронных сетей — достаточно упомянуть замечательную работу 2012 года. Krizhevsky и др. на турнире ImageNet Сhallenge. С тех пор, благодаря усовершенствованию архитектуры моделей, улучшению методов обучения, большим базам данных, например Open Images или ImageNet, и простым в использовании библиотекам, таким как TensorFlow и PyTorch, исследователи создали классификаторы изображений, способные распознавать тысячи категорий. Оцените, насколько качественно сегодня работает классификация изображений:

Классификация изображений хороша тем, что автоматически определяет объекты на изображении, но для эффективного поиска одной этой функции недостаточно. В самом деле, если пользователю нужно найти фото с пляжа, эта функция действительно может выдать фото с самыми высокими балльными оценками для этой категории, но что, если пользователь задаст поиск не пляжа, а берега? Что если пользователь задаст поиск не яблока, а фрукта или сорта яблок гренни смит?

Мы могли бы составить большой словарь синонимов и близких по значению слов, а также иерархических отношений между ними, но задача эта совершенно неподъёмная, особенно если вы поддерживаете несколько языков.

Векторы слов

Давайте несколько переформулируем задачу. Выход классификатора изображений можно интерпретировать как вектор j_c балльных оценок по каждой категории. Данный вектор определяет содержание изображения как точку в C-мерном пространстве категорий, где C — количество категорий (несколько тысяч). Если нам удастся извлечь осмысленное представление запроса в таком пространстве, мы сможем понять, насколько близко вектор изображения находится к вектору запроса, и принять это за меру соответствия изображения запросу.

Получение векторных представлений текста — предмет большого количества исследований в области обработки с использованием естественного языка. Один из наиболее известных методов в этой области — word2vec — описан в статье Mikolov и др. в 2013 году. Каждому слову в словаре функция Word2vec присваивает собственный вектор, и, исходя из этого, слова имеющие сходные значения, будут иметь близкие друг к другу векторы. Эти векторы находятся в d-мерном векторном пространстве слов, где d обычно составляет несколько сотен.

Узнать векторное представление поискового запроса можно, просто просмотрев его вектор word2vec. Этот вектор не является частью пространства категорий векторов классификатора изображений, однако можно преобразовать его в пространство категорий, ссылаясь на названия категорий изображений следующим образом:

1. Для слова запроса q можно получить d-мерный вектор слов q_w, нормализованный к единичному вектору. Для векторов в пространстве слов будем использовать нижний индекс w, а для векторов в пространстве категорий — нижний индекс c.

2. Для каждой категории получаем нормализованный вектор слов для названия категории cⁱ_w. Затем определяем значение m̂i = q_w - cⁱ_w — косинусное сходство векторов запроса и i-й категории. Полученная оценка от -1 до 1 показывает, насколько близко слово запроса соответствует названию категории. Сводя отрицательные значения к нулю (то есть применяя функцию mi = max(0, m̂i)), получаем оценку в том же диапазоне, что и результаты классификатора изображений.

3. Таким образом, можно вычислить q_c = [m1 m2 ... mC], вектор в C-мерном пространстве категорий, показывающий, насколько близко запрос соответствует каждой из категорий — аналогично тому, как вектор классификатора изображений для каждого изображения показывает, насколько близко изображение соответствует каждой из категорий.

Шаг 3 — осуществляем обычное векторно-матричное умножение, q_c = q_wC, где C — матрица со столбцами векторов слов категории cⁱ_w.

После сопоставления запроса с вектором пространства категорий q_c к вектору пространства категорий для каждого изображения можно применить косинусный коэффициент и получить таким образом окончательный балл релевантности для изображения — s = q_cj_c.

Мы получили функцию релевантности. С её помощью мы ранжируем изображения побалльно и формируем список результатов запроса. Применение данной функции к набору изображений можно также выразить как векторно-матричное умножение, s = q_cJ, в котором каждый столбец J представляет собой выходной вектор классификатора j_c для изображения, а s — вектор балльных оценок релевантности для всех изображений.

Пример

Рассмотрим пример с небольшим количеством измерений. В нём векторы слов имеют три измерения, а классификатор — четыре категории: яблоко, пляж, одеяло и собака.

Предположим, пользователь искал берег. Просматриваем вектор слов, получаем [0,35–0,62 0,70], затем умножаем на матрицу векторов слов категорий и проецируем запрос на пространство категорий.

Поскольку вектор слова берег близок к вектору слова пляж, данная проекция имеет высокое значение в категории пляж.

Описание модели

Наш классификатор изображений представляет собой  сеть EfficientNet, натренированную на наборе данных OpenImages. Классификатор выдаёт балльные оценки примерно по 8 500 категориям. Мы выяснили, что данная архитектура и набор данных обеспечивают хорошую точность при разумных затратах. Это довольно важно, если приходится работать с клиентской базой Dropbox.

Для тренировки и запуска модели будем использовать TensorFlow и  заранее обученные векторы слов ConceptNet Numberbatch. Полученные результаты оказались удовлетворительными, и, что важно для нас, модель способна работать с несколькими языками, возвращая сходные векторы для слов на разных языках со сходными значениями. Это упрощает процесс поиска изображений на разных языках: векторы слов dog на английском языке и chien на французском сходны, поэтому мы можем осуществлять поиск на обоих языках без перевода поискового запроса с одного языка на другой.

Для поиска по нескольким словам будем анализировать запрос как логическую операцию AND, применённую к отдельным словам. Мы также ведём перечень терминов, состоящих из нескольких слов, например beach ball, которые можно рассматривать как одно слово. Если запрос содержит один из таких терминов, мы делаем альтернативный синтаксический разбор и применяем операцию OR для двух разобранных запросов, после чего запрос beach ball принимает вид (beach AND ball) OR (beach ball). Этот термин будет подходить для поиска как больших разноцветных надувных мячей для игры на пляже, так и для теннисных мячей на песке.

Архитектура в производственной среде

Каждый раз, когда пользователь выполняет поиск, получать полную и актуальную матрицу J не вполне практично. Пользователь может иметь доступ к сотням тысяч или даже миллионам изображений, а наш классификатор способен работать с тысячами измерений, поэтому такая матрица может содержать миллиарды записей и должна обновляться каждый раз, когда пользователь добавляет, удаляет или изменяет изображение. Для сотен миллионов пользователей такой подход просто нереализуем (пока).

Поэтому вместо создания экземпляра J для каждого запроса мы используем описанный выше приближённый метод, который может быть эффективно реализован на поисковом движке Dropbox Nautilus.

По сути, Nautilus состоит из прямого индекса (forward index), сопоставляющего каждый файл с определёнными метаданными (например, именем файла) и полным текстом файла, и инвертированного индекса (inverted index), сопоставляющего каждое слово со списком публикаций (posting list) для всех файлов, содержащих это слово. При текстовом поиске содержимое индекса для нескольких файлов с рецептами может выглядеть примерно так:

Если пользователь ищет белое вино, мы ищем эти два слова в инвертированном индексе и обнаруживаем в doc_1 и doc_2 оба слова, поэтому мы должны включить их в результаты поиска. В doc_3 есть только одно из этих слов, поэтому мы должны либо опустить его, либо поставить последним в списке результатов. 

После того как найдены все документы, в которых могут встречаться устраивающие нас результаты, мы просматриваем их в прямом индексе и используем полученную таким образом информацию для ранжирования и фильтрации. В данном случае doc_1 можно поставить выше, чем doc_2, так как два искомых слова в doc_1 идут подряд.

Применение методов текстового поиска для поиска изображений

Систему текстового поиска можно использовать для реализации нашего алгоритма поиска изображений. В прямом индексе можно хранить вектор пространства категорий j_c каждого изображения. В инвертированном индексе для каждой категории будет храниться список публикаций изображений с положительными балльными оценками для данной категории.

Предположим, пользователь ищет слово пикник:

1. Находим вектор q_w для слова пикник и умножаем его на матрицу проекции пространства категорий C и получаем q_c, как описано выше. C — это фиксированная матрица, одинаковая для всех пользователей, поэтому её можно хранить в памяти.

2. Для каждой категории, имеющей ненулевую запись в q_c, получаем список идентификаторов из инвертированного индекса. Объединение таких списков представляет собой набор совпадающих изображений, полученных в результате поиска, однако эти результаты ещё необходимо ранжировать.

3. Для каждого результата поиска возьмём вектор пространства категорий j_c из прямого индекса и умножим его на q_c, чтобы получить балльную оценку релевантности s. Получаем результаты с балльной оценкой выше порогового значения, ранжированные по этой балльной оценке.

Оптимизация для масштабируемости

Описанный подход по-прежнему связан со значительными затратами как с точки зрения пространства для хранения данных, так и с точки зрения времени обработки запросов. Для 10 000 категорий в прямом индексе для каждого изображения необходимо хранить 10 000 балльных оценок классификатора, что при использовании четырёхбайтовых значений с плавающей запятой составляет 40 КБ. А поскольку оценки классификатора редко когда бывают строго нулевыми, в большинство из таких 10 000 списков идентификаторов будет добавлено стандартное изображение. Если для идентификаторов изображений использовать четырёхбайтовые целые числа, добавятся ещё 40 КБ и стоимость индексирования одного изображения составит 80 КБ. Другими словами, для многих изображений их индекс будет больше по размеру, чем сам файл изображения!

Что касается времени обработки запроса, под которым понимается задержка, с которой система выдаёт пользователю результаты поиска, мы можем ожидать, что около половины балльных оценок совпадения запроса и категории m̂i будут положительными, поэтому нам придётся прочитать около 5 000 списков идентификаторов из инвертированного индекса. Если сравнивать с текстовыми запросами, то они обычно считывают всего около 10 списков публикаций.

К счастью, существует множество близких к нулю значений, которые вполне можно отбросить, и это не повлияет качество поиска. Оценка релевантности для каждого изображения была представлена выше как s = q_cj_c, где q_c — балльная оценка соответствия между запросом и каждой из примерно 10 000 категорий, а j_c — балльные оценки приблизительно 10 000 категорий, полученные с помощью классификатора. Оба вектора в основном состоят из близких к нулю значений, которые не вносят практически никакого вклада в s. 

В нашем приближении мы определим значения всех записей q_c и j_c как равные нулю, кроме нескольких самых больших. Экспериментально было выяснено, что сохранение 10 записей в q_c и 50 записей в j_c вполне достаточно для получения не менее содержательного итогового результата. Экономия на хранении и обработке данных получилась довольно существенной:

• В прямом индексе вместо 10 000-мерных плотных векторов мы храним разрежённые векторы с 50 ненулевыми записями, то есть с 50 записями с наиболее высокими балльными оценками категорий для каждого изображения. В разрежённом представлении мы храним позицию и значение каждой ненулевой записи; 50 двухбайтовых позиций (записанных целыми числами) и 50 четырёхбайтовых значений (записанных в виде числа с плавающей точкой) требуют всего около 300 байт.

• В инвертированном индексе каждое изображение добавляется не в 10 000, а в 50 списков идентификаторов, что обходится всего в 200 байт. Таким образом, общий размер индекса для каждого изображения составляет 500 байт вместо 80КБ.

• Во время запроса q_c имеет 10 ненулевых записей, поэтому нужно просканировать только 10 списков идентификаторов — примерно такой же объём работы выполняется в текстовых запросах. Набор результатов сужается, и такой набор можно быстрее оценить.

Благодаря описанной выше процедуре оптимизации резко снижаются затраты на индексирование и хранение данных, а задержки выполнения запросов становятся сопоставимыми с задержками при текстовом поиске. Таким образом, если пользователь запускает операцию поиска, можно параллельно выполнять поиск по тексту и изображениям и выдавать полный набор результатов совокупно, причём пользователь не будет ждать результатов поиска изображений — они будут выведены одновременно с результатами поиска по тексту.

Доступность функций

Функция поиска в изображениях в настоящее время доступна для всех наших пользователей профессионального и бизнес-уровня. Объединяя функции поиска в любых изображениях, OCR-поиска в изображениях и документах и полнотекстового поиска в текстовых документах, мы делаем большинство файлов пользователей доступными для поиска на основе содержимого.

Поиск в видео?

Естественно, мы работаем над функцией поиска содержимого в видеофайлах, хранящихся в Dropbox. Создание функции поиска изображений приближает нас к решению этой задачи. Со временем, мы надеемся, пользователи смогут осуществлять поиск данных и в видеофайлах. Методы поиска одного кадра в видео или индексирования всего клипа для поиска, возможно, посредством совершенствования техник работы с неподвижными изображениями, пока ещё находятся на стадии разработки, но вспомните, что всего несколько лет назад фраза "найти все фотографии, где я на пикнике со своей собакой" работала только в голливудских фильмах.

Искусственный интеллект сегодня экономит не только время, но и силы, которые направляются на решение всё более сложных задач. И если вам интересно решать задачи в области искусственного интеллекта, вы можете обратить внимание на наш курс "Machine Learning и Deep Learning", разработанный совместно с компанией NVIDIA.

Узнайте, как прокачаться и в других специальностях или освоить их с нуля:

• Профессия Data Scientist

• Профессия Data Analyst

• Курс по Data Engineering