О себе

Здравствуй, Хабр! Меня зовут Павел, я работаю техническим директором в компании, занимающейся производством IoT устройств. Производим много чего — начиная от контроллеров для умных домов, заканчивая умными приборами учёта на своём запатентованном протоколе сенсорных сетей.

Также исполняют обязанности генерального директора ит-компании. В прошлом полуфиналист ЧМ по программированию ACM ICPC.

Мотивация

Пишу я это статью потому, что наша команда убила около месяца на поиск решения (ещё недели две на реализацию и написание тестов) для хранения и эффективного поиска распознанных лиц в базе данных, с целью сэкономить время вам в ваших проектах. Спойлер: ничего готового вроде классного плагина для существующей СУБД не нашли, а сроки полыхали, по этому написали свою СУБД именно для этой задачи (хранения огромного количества эмбендингов лиц). Моя статья ни в коем случае не претендует на звание исчерпывающего руководства, но, я надеюсь, что она даст точку старта для дальнейшего изучения и развития наших мыслей.

Эмбеддинг – это отображение из дискретного вектора категориальных признаков в непрерывный вектор с заранее заданной размерностью.

Итак, с чего всё начиналось

Как-то пришла нам задачка на разработку системы учёта уникальных, постоянных и редких посетителей бутиков в торговых центрах с множеством разных фич, таких как распознавание возраста посетителей, эмоций на входе / выходе из бутика, времени нахождения людей в ТЦ и прочие важные для арендодателей и арендаторов вещи. Мы, не долго думая, приступили к реализации, написали структуру, обучили каскад нейронных сетей для решения данных задач. "Честная" точность на выходе у нас получилась на тот момент 87%, и это, как кажется, неплохой результат с синтетически сгенерированными и раздутыми данными. Со временем продакшена слои отвечающие за те или иные фичи дообучались, работали разметчики чуть ли ни в 3 смены. Через год обкатали решение на порядка десяти мелких и средних торговых центрах в Подмосковье. Изначально это было облачный решением и мы работали с 2-3 собственниками сетей ТЦ. Как только стали расширяться и подключать другие города — начался настоящий ад… И дело не в растущей нагрузке на сервисы с нейронками, а в том, что кол-во "уникальных" (насколько позволяла точность нейронок) лиц перевалило за 6 миллионов. И вот тогда мы начали искать другие варианты решения задачи о хранении лиц.

Постановка задачи

Для тех, кто ещё не понял, повторяю: наша проблема (и тема данной статьи) была не в том, чтобы написать нейронки, не в том, чтобы написать систему аналитики и визуализации всего этого зоопарка аналитических данных, которые мы получаем на выхлопе, а в том, чтобы эффективно хранить данные о лицах. Запросы на запись у нас были не особо большими пачками в среднем в зависимости от дня недели и времени дня от 10 до 200 за один запрос при условии, что лицо мы видим первый раз. Если говорить о запросах на проверку лица в базе, то тут уже пачки доходили до 50к за раз. В среднем стабильно 10к.

Итак, требования к запросам для СУБД +- сформировали, приступим к рассмотрению эволюции наших мыслей для решения данной задачи.

Рассматривать решение на примере нашей интеллектуальной собственности я, увы, не могу по очевидным причинам, поэтому рассмотрим упрощённую задачу на примере простой опенсорсной распознавалки лиц от опенсорсной dlib, в которой вся работа происходит с помощью вызова трёх высокоуровневых функций. При этом она написана на современном C++, использует BLAS, имеет обёртку на Python и достаточно быстро работает на CPU. Мой выбор для демонстрации нашего исследования пал именно на неё.

Если мне не изменяет память, нейронка в dlib обучена на минимизации Евклидова расстояния между похожими эмбендингами, в доке рекомендуют ставить триггер на срабатывание поиска лица в базе на расстояние 0.6, от этого и будем отталкиваться. На выхлопе мы получаем эмбендинг размером 128.

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

Перейдём к вариантам решения.

Тупо перебор

Что тут говорить. Просто берём и с помощью dlib строим нужное количество эмбендингов и в цикле ищем минимальное расстояние для целевого.

def get_img_vector(img):
    dets = detector(img, 1)
    for k, d in enumerate(dets):
        shape = sp(img, d)
        return facerec.compute_face_descriptor(img, shape)
    return None

def prepare_database():
    database = {}

    for file in glob.glob("_images/*"):
        identity = os.path.splitext(os.path.basename(file))[0]
        img = cv2.imread(file, 1)

        database[identity] = get_img_vector(img)

    return database

def who_is_it(img, shape, database):
    face_descriptor1 = facerec.compute_face_descriptor(img, shape)

    min_dist = 100
    identity = None

    for (name, db_enc) in database.items():
        dist = distance.euclidean(db_enc, face_descriptor1)

        if dist < min_dist:
            min_dist = dist
            identity = name

    print(min_dist)
    if min_dist > 0.57:
        return None
    else:
        return str(identity)

if __name__ == "__main__":
    global sp
    sp = dlib.shape_predictor('weights/shape_predictor_face_landmarks.dat')
    global facerec
    facerec = dlib.face_recognition_model_v1('weights/dlib_face_recognition_resnet_model_v1.dat')
    global detector
    detector = dlib.get_frontal_face_detector()

    database = prepare_database()
    webcam_face_recognizer(database)

В функции webcam_face_recognizer происходит получение кадров с камеры в цикле (обычная cv2-шная тема) и при поиске лица получает его эмбендинг. Нас интересует функция who_is_it, в ней происходит полный проход по базе и поиск эмбендинга с минимальными расстоянием к только что полученному. Вот, в принципе, и всё, поздравляю, задача решена!

Очевидно, такой подход не примут даже если бы мы писали лабораторку на 1м курсе провинциального вуза. Сложность данного решения получилась О(N*k), где N — у нас количество лиц в базе, а k — размерность пространства поиска. Осталось найти побольше данных для базы и убедиться в том, что, долго это решение жить не будет. Все эти темы слишком объёмные, чтобы подробно рассказывать про них с огромными листингами кода, поэтому я надеюсь на ваше знакомство с задачей распознавания лиц хотя-бы на базовом уровне и просто изложу остальные этапы эволюции наших мыслей. Замеры даже делать бессмысленно.

Матрица

Думаете необходимо перебирать список всех эмбедингов, чтобы найти совпадение?

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

scores = np.linalg.norm(face_descriptor1 - np.asarray(database.items()), axis=1)
min_el_ind = scores.argmax()

Это уже большой буст в плане производительности, но, оказывается, можно ещё чуточку быстрее.

Поиск можно значительно ускорить, немного проиграв в его точности, используя библиотеку nmslib. Она использует метод HNSW для приближенного поиска k ближайших соседей. По всем имеющимся векторам должен быть построен так называемый индекс, в котором затем и будет производиться поиск. Создать и сохранить на диск индекс для евклидовой дистанции можно следующим образом:

import nmslib

index = nmslib.init(method='hnsw', space='l2', data_type=nmslib.DataType.DENSE_VECTOR)

for idx, emb in enumerate(database.items()):
    index.addDataPoint(idx, emb)

index_params = {
    'indexThreadQty': 5,
    'skip_optimized_index': 0,
    'post': 3,
    'delaunay_type': 2,
    'M': 100,
    'efConstruction': 2000
}

index.createIndex(index_params, print_progress=True)
index.saveIndex('./db/database.bin')

Подробнее о HNSW можно почитать в этой классной статье.

Но нам не нужна "примерная" точность в угоду производительности. Зря что ли старались наши МЛщики?

Кластеризация

Временное решение, которое проработал около 4х месяцев. Код писать под адаптацию dlib долго, по этому просто прикреплю классную визуализацию всего этого дала.

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

Плагин для postgresql

Нам не подошёл данный вариант решения из-за особенностей непростой формулы расстояния, на котором обучалась наша нейронка, но как вариант с Евклидовым расстоянием (косинусным, или ещё более простыми мат. формулами) стоит рассмотреть.

Инициализация:

import postgresql

def setup_db():
    db = postgresql.open('pq://user:pass@localhost:5434/db')
    db.execute("create extension if not exists cube;")
    db.execute("drop table if exists vectors")
    db.execute("create table vectors (id serial, file varchar, vec_low cube, vec_high cube);")
    db.execute("create index vectors_vec_idx on vectors (vec_low, vec_high);")

Вставка:

query = "INSERT INTO vectors (file, vec_low, vec_high) VALUES ('{}', CUBE(array[{}]), CUBE(array[{}]))".format(
            file_name,
            ','.join(str(s) for s in encodings[0][0:64]),
            ','.join(str(s) for s in encodings[0][64:128]),
        )
db.execute(query)

Поиск:

import time
import postgresql
import random

db = postgresql.open('pq://user:pass@localhost:5434/db')

for i in range(100):
    t = time.time()
    encodings = [random.random() for i in range(128)]

    threshold = 0.6
    query = "SELECT file FROM vectors WHERE sqrt(power(CUBE(array[{}]) <-> vec_low, 2) + power(CUBE(array[{}]) <-> vec_high, 2)) <= {} ".format(
        ','.join(str(s) for s in encodings[0:64]),
        ','.join(str(s) for s in encodings[64:128]),
        threshold,
    ) + \
            "ORDER BY sqrt(power(CUBE(array[{}]) <-> vec_low, 2) + power(CUBE(array[{}]) <-> vec_high, 2)) ASC LIMIT 1".format(
                ','.join(str(s) for s in encodings[0:64]),
                ','.join(str(s) for s in encodings[64:128]),
            )
    print(db.query(query))
    print('inset time', time.time() - t, 'ind', i)

Поиск в базе с 10^5 лиц на моём ноутбуке (4-х ядерный i5, 2,33 GHz) занимал 0.034 сек.
Как вам? По моему +- сказка для средних объёмов базы лиц (домофоны с распознаваним лиц точно потянет). Тут и известная СУБД с большим комьюнити, тут и опыт шардирования данной БД есть, и репликации… Слюнки текли, а мы двинулись дальше.

K-d дерево

Наконец-то дошли до, как нам казалось, "серебреной пули", которую мы используем и по сей день. Боже, почему мы не наткнулись на это сразу.

В общем, пару слов про эту, казалось бы, нехитрую структуру данных.

K-d дерево — структура данных с разбиением на пространства для упорядочивания точек в k-мерном пространстве. Такие деревья используются для некоторых приложений, таких как поиск в многомерном пространстве ключей (поиск диапазона, поиск ближайшего соседа и т. д.), ну и по сути своей представляют особый вид двоичных деревьев поиска.

Если хочется больше подробностей, прошу:

Подробнее тут.

Расход памяти в О(N) нам отлично подходил, но на начальном этапе, не имея практических замеров на реальных данных нас напрягали вставка, поиск и удаление в худшем случае за О(N). Но опасения оказались напрасны, на практике было всего прекрасно!

Вот исходный код нашего "кастрированного" дерева, дабы соблюсти NDA.

На этом мы и остановись. Месяц и две недели ушли не зря, теперь мы можем спать спокойно.

На данный момент

Сейчас наше решение отлично работает с хранением 15 миллионов уникальных лиц. Это количество меняется, т. к. МЛщики выкатывают новые веса и графы для неронок и уникальность приходится пересчитывать.

Также я уже упоминал потенциальную зону роста нашего решения — автоматическая балансировка при вставке. Учитывая, что сейчас не лучшее время для посещения ТЦ в связи с эпедемиологической обстановкой, у нас есть время для эээээксперементов :)

Наше k-tree деградирует в разы медленнее, чем решение с кластеризацией и перебалансировку можно проводить раз в пол года учитывая нашу нагрузку (особенно в данный период), но мы же параноики перфекционисты, перебалансировка у нас происходит раз в неделю (ночь вскр-пн), на каждой реплике операция перебалансировки занимает около 4-6 секунд. Поэтапная перебалансировка всех реплик занимает чуть больше минуты.

Подводя итоги

Многие могут сказать, мл можно было заюзать шардирование БД и всё было бы ок. Отвечу — за 1,5 недели до предполагаемого момента когда мы начинали бы захлёбываться, нам бы не удалось реализовать шардирование, ни на уровне приложения, возможно, только на уровне СУБД, но я считаю, что этим мы бы только отсрочили неизбежное. Можно было бы разделить сервис и запустить копию для каждого ТЦ отдельно, но этим бы мы не ответили на некоторые вопросы, которые интересуют наших заказчиков.

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

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

А мораль сей басни такова — толпою валят даже льва опенсорсные решения, испытанные временем. Учите алгоритмы, пасаны :)