- PVSM.RU - https://www.pvsm.ru -
В user-generated проектах часто приходится бороться с дубликатами, а для нас это особенно актуально, так как основной контент мобильного приложения iFunny — это изображения, которые постятся десятками тысяч ежедневно. Для поиска повторов мы написали отдельную систему, чтобы облегчить процесс и сэкономить море времени.
Под катом рассмотрим используемые инструменты, а потом перейдём к примеру реализации.
Существует огромное количество различных алгоритмов поисков дубликатов, каждый со своими плюсами и минусами. Один из таких — поиск наиболее похожих (близких) векторов, полученных с помощью CNN-сетей [1].
После классификации изображения через CNN-сеть на выходе получается вектор «того, что увидела сеть на изображении». В теории такой способ должен быть менее чувствительным к кропу, но ложных похожих изображений будет больше в сравнении с более точными методами.
Есть и другой недостаток. На выходе классификации получается большой вектор (2048 float для resnet152), который где-то нужно хранить и иметь возможность за разумный промежуток времени найти все N похожих векторов для искомого — что само по себе уже непросто.
Поиск наиболее близких векторов — частая задача, для решения которой уже есть отличные инструменты. Здесь лидером считается библиотека FAISS [2] от Facebook. Она использует эффективную кластеризацию векторов, позволяя организовывать поиск даже для векторов, которые не помещаются в RAM.
Но с FAISS напрямую работать не очень удобно. Это не база данных, нельзя туда просто сохранить вектор и запросить похожий (к тому же, после создания индекса его можно только пересоздать). Поэтому для промышленной эксплуатации нужно строить свою обвязку вокруг системы индексации.
Для этого есть весьма перспективный проект Milvus [3], который по дизайну сильно напоминает Elasticsearch. Отличие только в том, что Elasticsearch построен вокруг индекса lucene, а в Milvus вся архитектура выстроена вокруг индекса FAISS.
Структура коллекций тоже схожа:
Для каждой коллекции можно создать несколько партиций, по которым потом будет ограничиваться поиск. Партиция состоит из сегментов, которые представляют из себя простой набор файлов с id, исходными индексами и служебными данными.
Информация о коллекциях, партициях и сегментах хранится в отдельной SQL-базе. Для standalone-запуска используется встроенный SQLite, а ещё есть возможность использовать внешнюю MySQL-базу.
Проект Milvus находится в активной разработке (текущая версия 0.11.0). Пока что в нём нет репликации данных, как и возможности использовать другие SQL (или NoSQL) базы в качестве хранилища метаинформации. Поэтому пока для HA-решений можно использовать только схему с двумя экземплярами с общим хранилищем: один будет запущен, а другой — «спать». Для масштабирования можно будет использовать Mishards [4], но в 0.11.0 он сломан.
Кроме того, в 0.11.0 появилась возможность вместе с самим вектором и id сохранять в коллекцию дополнительные данные. Правда, пока без дополнительных индексов для них, но с возможностью поиска.
С точки зрения использования, Milvus выглядит как обычная внешняя база данных. Есть API (gRPC-клиент и набор http-методов) для сохранения и поиска вектора, управления коллекциями и индексами, а также для получения информации обо всех сущностях.
При создании коллекции можно указать максимальное количество записей в сегменте (segment_row_limit). Если превысить этот лимит, то Milvus начнёт строить индекс FAISS. С этим связана одна из особенностей Milvus: для всех добавляемых векторов, по которым ещё не создан индекс, поиск будет работать на основе полного перебора. Поэтому при больших значениях segment_row_limit будет много записей, для которых индекс ещё не построен (он также влияет ещё и на то, сколько будет созданных сегментов для коллекции). Для поиска похожих векторов в коллекции необходимо сделать поиск в каждом сегменте — и чем их больше, тем дольше поиск.
Обратите внимание, новый созданный сегмент не набивается до лимита при добавлении новых записей. Вместо этого происходит постепенное объединение сегментов по принципу игры 2048 [5] (до тех пор пока размер не превысит лимит). Поэтому при указании большего значения segment_row_limit может оказаться много меньших сегментов, для которых нет индекса, а значит поиск по ним будет медленным.
Несмотря на все особенности, поиск по векторам работает быстро. Архитектура индексов FAISS и самого Milvus позволяет за раз искать одновременно значения по нескольким векторам. И на практике последовательный поиск двух векторов будет существенно медленнее поиска обоих векторов за раз [6].
Milvus можно запускать как в CPU-версии, так и в GPU. Первую лучше всего использовать на процессорах, которые поддерживают инструкцию AVX512 [7]. Для этого достаточно просто запустить контейнер:
docker run -d --rm --name milvusdb -p 19530:19530 -p 19121:19121
milvusdb/milvus:0.11.0-cpu-d101620-4c44c0
В данном случае 19530 будет портом для gRPC-клиента, а 19121 — для http API [8].
В качестве CNN-сети можно взять любую из предобученных (или обучить самому) — результаты могут немного отличаться. В данном примере будем использовать предобученный resnet152:
model = models.resnet152(pretrained=True)
Вектор будем снимать со слоя `avgpool`:
layer = model._modules.get('avgpool')
А сам вектор получать с помощью hook:
vector = torch.zeros(2048)
def copy_data(m, i, o):
vector.copy_(torch.reshape(o.data, (1, 2048))[0])
hook = layer.register_forward_hook(copy_data)
model(prepared_image)
hook.remove()
Полный код получения вектора выглядит так:
import numpy as np
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
from PIL import Image
model = models.resnet152(pretrained=True)
layer = model._modules.get('avgpool')
model.eval()
pipeline = [
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]
def _prepare_Image(img: Image) -> Variable:
raw = img
for action in pipeline:
raw = action(raw)
return Variable(raw.unsqueeze(0))
def image_vectorization(image_path: str) -> np.ndarray:
img = Image.open(image_path)
prepared_image = _prepare_Image(img)
vector = torch.zeros(2048)
def copy_data(m, i, o):
vector.copy_(torch.reshape(o.data, (1, 2048))[0])
hook = layer.register_forward_hook(copy_data)
model(prepared_image)
hook.remove()
# vector normalization
norm_vector = vector / torch.norm(vector)
return np.array(norm_vector)
Теперь понадобится клиент для работы с Milvus. Можно взять любой из поддерживаемых (Python, Java, Go, Rest, C++). Возьмём Java-клиент и напишем пример на Kotlin. Почему? А почему бы и нет.
Подключаем Milvus SDK:
implementation("io.milvus:milvus-sdk-java:0.9.0")
Создаём подключение к Milvus:
val connectParam = ConnectParam.Builder()
.withHost("localhost")
.withPort(19530)
.build()
val client = MilvusGrpcClient(connectParam)
Создаём коллекцию под 2048 вектор:
val collectionMapping = CollectionMapping.create(collectionName)
.addVectorField("float_vec", DataType.VECTOR_FLOAT, 2048)
//выключаем автосоздание id
.setParamsInJson(JsonBuilder()
.param("auto_id", false)
.param("segment_row_limit", segmentRowLimit)
.build()
)
client.createCollection(collectionMapping)
Создаём IVF_SQ8 индекс:
Index.create(collectionName, "float_vec")
.setIndexType(IndexType.IVF_SQ8)
.setMetricType(MetricType.L2)
.setParamsInJson(JsonBuilder()
.param("nlist", 16384)
.build()
)
client.createIndex(index)
Сохраняем несколько векторов в коллекцию:
InsertParam.create(collectionName)
.setEntityIds(listOf(1L, 2L))
.addVectorField("float_vec", DataType.VECTOR_FLOAT, listOf(vector1, vector2))
client.insert(insertParam)
client.flush(collectionName) // чтобы сразу можно было найти вектор
Ищем ранее сохранённый вектор:
val dsl = JsonBuilder().param(
"bool", mapOf(
"must" to listOf(
mapOf(
"vector" to mapOf(
"float_vec" to
mapOf(
"topk" to 10,
"metric_type" to MetricType.L2,
"type" to "float",
"query" to listOf(vector1),
"params" to mapOf("nprobe" to 50)
)
)
)
)
)
).build()
val searchParam = SearchParam.create(collectionName)
.setDsl(dsl)
val result = client.search(searchParam)
println(result.queryResultsList[0].map { it.entityId to it.distance })
Если всё работает и правильно настроено, то вернётся похожий результат:
[(1, 0.0), (2, 0.2)]
Для первого вектора расстояние L2 с самим собой будет 0, а с другим вектором — больше 0.
Всё вышеприведённое, конечно, только наброски, но этого достаточно, чтобы попробовать создать Python-сервис для классификации и получения вектора. И либо для него накрутить API для сохранения и поиска векторов, либо сделать в отдельном сервисе (например, на Kotlin), который будет получать вектор и сохранять его уже в Milvus самостоятельно.
Спасибо всем, кто дочитал до конца, надеюсь, вы нашли для себя что-то новое. А если вас заинтересовал проект Milvus, то можете поддержать его на Github [9].
Автор: Евгений Захаров
Источник [10]
Сайт-источник PVSM.RU: https://www.pvsm.ru
Путь до страницы источника: https://www.pvsm.ru/java/359969
Ссылки в тексте:
[1] CNN-сетей: https://ru.wikipedia.org/wiki/%D0%A1%D0%B2%D1%91%D1%80%D1%82%D0%BE%D1%87%D0%BD%D0%B0%D1%8F_%D0%BD%D0%B5%D0%B9%D1%80%D0%BE%D0%BD%D0%BD%D0%B0%D1%8F_%D1%81%D0%B5%D1%82%D1%8C
[2] FAISS: https://github.com/facebookresearch/faiss
[3] Milvus: https://www.milvus.io/
[4] Mishards: https://github.com/milvus-io/milvus/tree/master/shards
[5] 2048: https://www.programmersought.com/article/93384828799/
[6] поиска обоих векторов за раз: https://github.com/milvus-io/milvus/issues/2437
[7] AVX512: https://medium.com/unstructured-data-service/milvus-performance-on-avx-512-vs-on-avx2-5dfb08d63a4c
[8] http API: https://github.com/milvus-io/milvus/blob/master/core/src/server/web_impl/README.md
[9] Github: https://github.com/milvus-io/milvus/
[10] Источник: https://habr.com/ru/post/534462/?utm_source=habrahabr&utm_medium=rss&utm_campaign=534462
Нажмите здесь для печати.