Как научиться распознавать новый документ за две недели?
Всем привет! Мы — команда распознавания документов Домклик, и мы помогаем банку упростить и ускорить процесс выдачи ипотеки.
Как думаете, сколько документов вам нужно отправить в банк, чтобы получить ипотеку? Если вы счастливчик, покупающий новостройку и уже являетесь зарплатным клиентом Сбера, то процедура максимально упрощена: никаких документов загружать не требуется! Но бывают и более сложные случаи, когда просят подтвердить доход с помощью 2-НДФЛ, или загрузить предыдущий договор купли-продажи, если покупаете вторичку. И в зависимости от ситуации может понадобиться еще много разных документов — в нашей базе более 200 позиций.
Для обработки всех документов мы используем свой алгоритм автоматического распознавания текста. Но с таким количеством разнообразных типов встает вопрос не только качества распознавания, но и скорости добавления и поддержки новых документов. Сегодня мы расскажем о нашей инфраструктуре, которая позволила всего за пару недель научиться распознавать свидетельство о заключении брака.
Архитектура распознавания документов
Что помогает нам быстро внедрить новый документ в систему и обучить необходимые модели? Секрет в том, что однажды мы сели и тщательно продумали систему обучения для каждого нового документа, а затем создали универсальный пайплайн для обучения моделей распознавания, каждая часть которого представляет собой отдельную библиотеку с инфраструктурой для обучения. Таким образом, для обучения новых моделей нам достаточно было создать датасет, выставить параметры обучения и запустить эксперимент!
Все загруженные в систему документы, в том числе свидетельство о браке, обрабатываются по одной схеме:
Сначала документ отправляется в классификатор, который определяет тип документа. Далее детектор дефектов проверяет, есть ли на фотографии сильные блики, находится ли документ в фокусе, не сильно ли обрезаны края. При необходимости мы попросим клиента переснять документ. Следующим шагом документы попадают на этап распознавания.
Пайплайн распознавания распознавания состоит из четырех частей:
определение угла поворота документа;
выделение полей, необходимых для распознавания;
распознавание текста;
оценка надежности прогноза.
А теперь чуть подробнее о каждом этапе. Клиенты загружают в систему документы в совершенно разной ориентации, а значит нам необходимо поворачивать снимки перед распознаванием полей. На фотографии документ располагается, в основном, вдоль границ изображения, поэтому нам достаточно легковесной сети для обучения классификатора на четыре класса: четыре поворота на 90 градусов, чтобы добиться высокого качества на тестовой выборке.
Далее документ передается в экстрактор полей. Чтобы увеличить качество распознавания текста, мы не используем решение для поиска всего текста на изображении, а предварительно выделяем необходимые для распознавания поля, чтобы передать модели OCR только кропы (фрагменты) изображения, на которых необходимо выделить текст нужного класса. О работе модели распознавания мы расскажем далее.
После того, как модель OCR распознает текст, необходимо оценить надежность этого прогноза. На выходе из экстрактора полей и модели распознавания мы получаем несколько различных оценок, которые можно интерпретировать как уверенность результата. Сначала мы объединяем их, чтобы получить одну величину, характеризующую уверенность пайплайна в распознавании конкретного поля, а затем подбираем пороговое значение, чтобы отнести предсказания модели к надежным или требующим дополнительной проверки.
Чтобы стандартизировать поток данных через нашу инфраструктуру, мы реализовали отдельную библиотеку. Она используется для хранения конфигурации, по которой строится уникальное внутреннее представление документа. Такое представление используется для хранения результатов каждой модели и может поэтапно наполняться углом поворота, затем координатами полей и текстом, и в конце — уверенностью. Эти данные потом можно использовать для обучения последующих этапов пайплайна или для анализа метрик.
ML-инфраструктура
Чтобы быстро получать готовую модель распознавания новых документов, мало просто иметь фреймворк для обучения. Еще необходимо организовать удобное хранение данных и логгирование результатов обучений.
Мы часто пополняем наш датасет новыми типами документов, а также проводим много экспериментов. Сохранять артефакты моделей и датасет на диск в отдельную директорию удобно лишь на начальной стадии разработки моделей, а при большой вариативности документов и моделей такая организация работы становится неэффективной.
Для решения этих проблем в ДомКлике создали единый инструмент для версионирования данных и логгирования результатов обучений — MLECO: Machine Learning Ecosystem. Он помогает ds-ам хранить в одном месте различные версии датасетов, а также параметры экспериментов, которые были запущены на этих датасетах в одном месте. Это позволяет вернуться к предыдущим обучениям, посмотреть, какие были метрики в конкретном эксперименте, и понять, был ли он удачным.
Модель распознавания: технические подробности
Давайте перейдем от инфраструктуры к техническим подробностям, которые позволили нам быстро и с хорошим качеством обрабатывать свидетельства о заключении брака. Поговорим о модели распознавания текста — ключевом компоненте всего пайплайна.
Общий датасет
Как известно, для обучения уникальных OCR-моделей с нуля требуется много размеченных данных. За несколько лет у нас скопился большой объем таких документов. Среди них можно встретить развороты паспортов разных стран, трудовые книжки, справки разных форм и стандартов. Документы встречаются как печатные, так и рукописные. Мы собрали один большой набор из фрагментов, содержащий миллион слов и фраз, на которых можно обучить универсальную базовую модель, подходящую для извлечения текста из любого документа.
Дообучение
У единой модели есть один главный недостаток — она очень долго обучается. Но имея на руках уже обученную модель всегда можно дообучить ее на новом документе. Для этого нам было достаточно добавить в обучающую выборку 1000 свидетельств о заключении брака.
Поскольку датасет для дообучения значительно меньше по объему, чем наш основной датасет, обучение заняло не дни, а часы при сопоставимом количестве эпох. Мы смогли сократить длительность обучения в 8 раз, со 120 до 15–18 часов. Также нам хватило 1000 документов и для обучения экстрактора полей, мы используем каскадную модель пакета Detectron2 с предобученными на ImageNet весами.
Стоит понимать, что в общем случае условные 1000 документов не надут нам идеального качества, поскольку влияет сложность документов, их вариативность, качество сканов и разметки. Но свидетельство о браке — достаточно простой и структурированный документ, поэтому мы быстро смогли получить рабочее решение с высокой долей корректно распознанных документов.
Сэмплирование кропов в батче
Поскольку в датасете для распознавания текста присутствуют разные документы, то среди кропов встречаются как короткие, так и длинные поля. При формировании одного батча из нескольких фрагментов разница в их длине компенсируется добавлением отступов (padding). А так как поля могут очень сильно различаться по длине, то и величина отступов может быть очень большой. При этом их количество растет по мере увеличения размера батча. А чем больше отступов, тем менее эффективно мы используем ресурсы видеокарты и медленнее обучаем. Чтобы побороть это, мы решили написать свой сэмплер.
Идея достаточно простая и состоит в том, чтобы в батч попадали кропы одинаковой или схожей длины. Чтобы этого добиться, для начала заранее сортируем и группируем все индексы по ширине кропа. Затем дважды перемешиваем: внутри каждой группы и после разделения на батчи. Такой подход, несмотря на большое количество данных в датасете ускорил обучение примерно на 40%.
Примеры кропов разного размера в батче
Обогащение датасета для обучения
На проде в модель распознавания будут приходить не идеальные, а полученные на предыдущем шаге экстрактора поля. Поэтому в дополнение к реальным регионам из разметки мы добавили в датасет регионы, полученные от экстрактора полей, обученного на предыдущем шаге пайплайна. Это позволило модели OCR лучше обучаться с учетом особенностей работы модели экстрактора, поскольку может меняться расположение региона на странице, его окружение и фон, а также варьируется ширина выделения кропа и эти особенности модели экстрактора сложно реализовать аугментацией.
Также обучение распознавания проводилось на синтетически сгенерированных данных в дополнение к реальным. Обычно на синтетике проводилось разогревание весов модели на начальных эпохах, а затем доля искусственных примеров в обучающем датасете уменьшалась на каждой эпохе.
LMDB-базы, сокращение времени на эксперименты
При объеме данных около 100 тыс. изображений и больше 1 млн фрагментов с текстом на все типы документов, перед нами встала задача эффективного хранения и использования общего датасета. Чтобы при обучении не тратить время на вырезание и масштабирование кропов во время подготовки батча, мы решили генерировать датасет с готовыми кропами перед обучением. Он сохраняется на диск и заново используется на каждой эпохе в текущем и во всех следующих обучениях.
Поскольку данных слишком много, мы не можем перед обучением загрузить все фрагменты в оперативную память. А так как чтение картинки и преобразование в тензор затратно по времени, то считывать с диска jpeg во время подготовки батча — тоже плохой вариант. Поэтому мы решили использовать для хранения LMDB (Lightning Memory-Mapped Database). Это быстрая по чтению и эффективная по памяти база данных в виде key-value хранилища. Она использует в качестве механизма работы с диском отображенные в память файлы (memory-mapped files), и благодаря этому кеширование данных выполняется на уровне ядра OC.
В промежуточной базе для каждого кропа мы храним все данные, которые требуются для формирования батча, то есть трехмерный тензор, текстовую строку, тип документа и поля, а также числовые значения для пространственных аугментаций. Поскольку в LMDB нет возможности фильтровать, то для этой задачи мы храним параллельный pandas.Dataframe.
Архитектурные нюансы модели
Архитектура модели распознавания устроена довольно просто. На вход подается изображение с текстом. Из региона с текстом сначала извлекаются визуальные признаки, а из них с помощью нескольких сверточных слоев мы получаем контекстуальные признаки. Для обучения всей модели используем ctc-loss (механика работы неплохо описана тут).
В качестве сети для извлечения признаков взяли сеть на основе ResNet. Мы модифицировали исходную архитектуру на основе этой статьи. Перепробовали много разных улучшений, но лучше всего сработало добавление BigLittle Module и увеличение количества BasicBlocks в residual-блоке. Дополнительно в BasicBlocks добавили SE-блоки для агрегации информации между каналами и увеличения качества работы сети. Для получения итогового результата — строки с текстом — используем beam search-декодирование с определенным алфавитом для каждого поля.
Квантизация и TorchScript
Итоговая модель для распознавания всех полей получилась довольно тяжелой и медленной. Чтобы ее облегчить, мы решили воспользоваться механизмом квантизации, реализованным в Pytorch. Он позволяет уменьшить размер обученной сети и ускорить инференс на продовых мощностях (в нашем случае серверы на CPU). Основная идея заключается в вычислении и хранении параметров сети с меньшей точностью: в частности, в виде int8 вместо чисел с плавающей точкой.
В проекте мы используем статическую квантизацию, поскольку динамическая не дает значимого ускорения для нашей архитектуры, а обучение с квантизацией слишком сильно замедляет этап обучения. На практике уменьшение точности после квантизации было незначительным (в районе 0,3%). При этом размеры весов уменьшились в четыре раза (сейчас 3 Мб), а скорость выросла в 3,5 раза.
Кроме этого, для сохранения квантизированных весов мы используем TorchScript, который позволяет сериализовать объект модели вместе со структурой классов (то есть вместе с описанием модели). Таким образом, после десериализации можно работать с моделью без импорта классов модели. Сериализация и десериализация выполняется одной строчкой. Это дало дополнительное ускорение на 10%.
Pytorch предоставляет три подхода к квантизации:
Динамическая квантизация. Помимо преобразования весов к int8 для выполнения вычислений также преобразует на лету активации в int8 и обратно. Таким образом, вычисления будут выполняться с использованием реализаций умножения и свертки матриц, эффективных на целочисленных значениях малой размерности, что приведет к более быстрым вычислениям. Нет значимого ускорения для нашей архитектуры.
Статическая квантизация — понижение точности весов уже обученной модели. При этом подходе выполняется дополнительный этап калибровки параметров квантизации с помощью подачи данных через сеть и вычисления результирующих распределений активаций. В частности, это делается с помощью вставки в разные точки модулей «наблюдателей», которые записывают эти распределения. Информация, накопленная наблюдателями, используется для определения параметров квантизации, которые следует использовать.
Обучение с квантизацией. При этом подходе все веса и активации «поддельно квантуются» во время как прямого, так и обратного проходов обучения: то есть значения с плавающей запятой округляются до целых значений, но все вычисления по-прежнему выполняются с числами с плавающей запятой. Таким образом, все корректировки весов во время обучения производятся с учетом того факта, что модель в конечном итоге будет квантована. Поэтому после квантования этот метод обычно дает более высокую точность, чем два других.
Подробнее про типы квантизации и их ограничения можно почитать здесь.
Подведение итогов
Сегодня мы рассказали об инфраструктурных особенностях нашей системы распознавания документов: пайплайне моделей, обогащении датасета данными, эффективном хранении и использовании этих данных, а также об архитектуре нейросети и квантизации. Одни из этих трюков позволяют ускорить обучение и инференс, другие — повысить качество распознавания, третьи — сократить объем разметки, требующейся для обучения моделей.
Вместе эти компоненты составляют достаточно сложную систему, но она позволяет быстро и предсказуемо вводить поддержку новых типов документов и экспериментировать с новыми подходами, улучшающими качество распознавания сразу на всех типах документов. И подводя итоги хочется отметить, что всего за две недели мы научились распознавать свидетельства о заключении брака с точностью 85+.
У нас большие планы на будущее, и если вам захотелось поучаствовать в наших проектах — будем вам рады.
