ML Tale: создание автоматизированного конвейера обучения машинному обучению и масштабирование сервиса

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

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

Архитектура

На этой диаграмме показана архитектура этого конвейера обучения и развертывания:

В левой части диаграммы показаны этапы обучения моделей с помощью AWS Batch.
Правая часть показывает, как мы используем эти модели для вывода в производстве.

В оставшейся части этой статьи мы подробно объясним каждый компонент.

Проблема

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

Это решение OCR превосходит по точности все другие протестированные нами решения. Однако экспериментальные модели машинного обучения (на самом деле 3 разные модели) были обучены специально для одного типа документов (водительские права Калифорнии, США).

Хуже того, общий размер этих трех моделей составлял около 300 МБ (к счастью, мы также уменьшили этот размер позже!). Чтобы применить такой уровень точности в производстве, нам нужно было обучить и распространить это на все доступные типы юридических документов в мире (около 6000 типов и версий).

Подводя итог, мы столкнулись с двумя проблемами:

Увеличьте охват, обучите и создайте больше моделей (столько же типов документов по всему миру)
Развертывание и масштабирование этого большого количества моделей в производстве

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

Автоматизация обучающего конвейера

Первой задачей было ускорить процесс обучения / выпуска новой модели для новых типов / версий документов. Изначально обучение одной новой модели заняло у нас около недели. Большую часть процесса пришлось выполнить вручную:

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

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

Как мы это сделали?

Инфраструктура: сначала мы поместили наш обучающий код в контейнер. Это помогло нам обеспечить предсказуемость при обучении новой модели, независимо от того, где и на какой платформе она была обучена (это также полезно для отладки). Более того, это дает нам хорошие возможности для использования облачных сервисов - в данном случае AWS Batch.
Автоматизация: чтобы решить проблему автоматизации, мы разделили каждую часть процесса обучения на разные логические задачи, которые можно запускать, тестировать и проверять независимо. Затем конвейер становится набором связанных задач, которые будут выполняться от начала до конца. В этом есть два преимущества: 1) если задача не выполняется из-за ошибки или проблемы с данными, например, мы можем перезапустить задачу после устранения проблемы и не будем повторно запускать уже выполненные задачи; 2) мы можем определить валидацию для каждой задачи, так что если, например, модель недостаточно хороша (то есть, она не проходит порог для проверки теста), мы не будем двигаться дальше, следовательно, не будем тратить ресурсы и Экономить деньги.

Инфраструктура

Чтобы проводить и масштабировать обучение в автоматическом режиме, нам нужна надежная инфраструктура. Мы использовали AWS Batch, который помог нам легко запустить несколько пакетов учебных заданий на AWS. Он динамически предоставляет запрошенные вычислительные ресурсы (например, экземпляры ЦП или ГП) на основе конкретных требований к ресурсам отправленных пакетных заданий. В Onfido мы используем Terraform для создания всей инфраструктуры в виде кода. На следующем рисунке показан снимок файлов Terraform, используемых для создания инфраструктуры. Вы можете найти код здесь.

Для каждого сервиса AWS создается файл terraform. На верхнем уровне нам потребовалось 3 сервиса AWS: пакетная обработка, S3 и реестр контейнеров (ECR).

Batch состоит из 3 частей: 1) Compute, где определены детали экземпляра ec2, 2) job, который содержит информацию об определенных заданиях, 3) Queue где он хранится на нескольких отправленных вакансиях.

iam.tf определяет рабак между службами. ecr.tf определяет репозиторий реестра контейнеров, в котором хранятся все образы докеров.

Мы также модулировали файл terraform, чтобы мы могли повторно использовать его для создания нового стека для других обучающих конвейеров в будущем. Другими словами, полностью новый стек инфраструктуры (со всеми службами и разрешениями) можно создать с нуля за несколько минут. Если кто-то хочет создать новую инфраструктуру, нам просто нужно заполните несколько параметров, как показано ниже.

module "prod_document_extraction" {
  source = "modules/training_stack"
  name   = "document_extraction"
  
  environment_resources = {
    memory = 240000
    vcpus  = 30
  }

  configuration = {
    instance_type = "p3.8xlarge"
    max_vcpus     = 128
    min_vcpus     = 0
    type          = "EC2" # SPOT
    image_id      = "ami-82defcfb"
  }

  # Add or remove env used in code
  environment_definition = <<EOF
  "environment": [
      { "name": "PYTHONPATH", "value": "." },
      { "name": "TRAINING_ID", "value": "<training id>" },
      { "name": "TRAINING_PATH", "value": "DATASETS/train" },
      { "name": "TEST_PATH", "value": "DATASETS/test" },
      { "name": "NUM_STEPS", "value": "50000" },
      { "name": "INTERVAL_SEC", "value": "1800" },
      { "name": "NUMBER_GPUS", "value": "4" },
    ]
  EOF
}

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

Автоматизация: создание цикла обучения

Мы разработали конвейер, чтобы упростить обучение нового документа. Конвейер разбивается на несколько связанных друг с другом задач. Каждая задача автономна и зависит от предыдущей задачи. Затем задачи выполняются последовательно одна за другой. Конечным результатом является сквозной конвейер, в который мы передаем несколько параметров, связанных с новым документом для обучения (тип документа, страна и т. Д.), Выполняем обучение и через несколько часов находим окончательные модели в S3.

В поисках возможных библиотек Luigi казался простой библиотекой, которая упрощает определение цепочки задач и сохраняет зависимости между ними. Это помогает определить зависимости между задачами, передать вывод задачи другой (если задача не удалась, мы можем перезапустить ее с того места, где она была неудачна, вместо того, чтобы заново запускать все заново).

Наш конвейер был немного сложнее, так как нам пришлось сгенерировать 3 модели с взаимозависимостями. Когда первая модель обучена, она будет использоваться для обучения следующей модели и так далее. Важно, чтобы точность каждой модели была приемлемой, потому что обученная модель будет использоваться для генерации входных данных второй модели. Мы определили и добавили задачу проверки после обучения каждой модели. Это поможет нам остановить обучение в любой момент, если точность модели не соответствует нашим требованиям (общий вариант использования: если данные плохо аннотированы).

Архитектура обучающего конвейера

Как уже упоминалось, мы стремились ускорить обучение за счет автоматизации ручных задач и возможности запускать несколько пакетов обучения параллельно . Каждое обучение не зависит от другого. Поэтому мы присвоили уникальный training_id(TID) каждому запуску и загрузили все связанные файлы и артефакты в S3 в соответствии с определенным training_id(TID). Мы использовали AWS Batch, чтобы превратить каждое обучение в самостоятельную работу. Поэтому мы можем расширить обучение, запустив их параллельно. Задание, выполняющее одного Луиджи, содержит связанные задачи.

Мы разделили весь конвейер на 3 большие категории, каждая из которых содержит несколько связанных задач. (Pre_Train → Train_Model → Post_Train)

Подзадачи Pre_Train:

Передайте соответствующие параметры (это создаст ветку git и сгенерирует код конфигурации. Мы используем jinja для генерации необходимого кода из шаблона.)
Запрос и извлечение релевантных аннотаций
Создайте обучающий и тестовый набор (изображение + аннотации): возьмите аннотации + изображения → чистый набор данных → разделите данные на обучающий и тестовый набор.

Любой, кто работал с данными, знает, что данные никогда не бывают чистыми. Нам пришлось принять различные меры предосторожности, чтобы найти проблемные изображения / аннотации. Этот шаг снизил вероятность сбоя и прерывания тренировочного цикла, что дорого с точки зрения денег и времени.

4. Загрузите указанные выше наборы в корзину S3 по созданному пути training_id(TID).

Подзадачи Train_Model:

Модель поезда 1
Подтвердите Model1 (если да → перейдите к 3, иначе уведомить и завершить)
Модель поезда 2
Подтвердите Model2 (если да → перейдите к 3, иначе уведомить и завершить)
Модель поезда 3
проверить модель3 (если да → перейти к Post Train else notify & end)

Подзадачи Post_Train:

Провести финальную проверку точности
создать репрезентативную оценку для этого документа (X%)
проверить порог X ≥
(если доволен) развернуть службу

Масштабирование развертывания / обслуживания

Горизонтальное и вертикальное масштабирование

Как уже упоминалось, размер созданной модели для каждого документа составлял около 300 МБ. Во время логического вывода модель будет загружена в память при инициализации службы, а затем она будет готова к обслуживанию запросов. Вначале сервис был стабильным и успешно обслуживал запросы. На данный момент у нас был один набор / сервис развертывания / реплик. Однако по мере того, как процесс обучения стал автоматизированным, мы начали вводить больше документов. Это постепенно привело нас к трем проблемам:

Образ докера рос в размерах: раньше мы запекали модели в образе и развертывали сервис в Kubernates.
Нам нужно было выделить сервису больше оперативной памяти, чтобы иметь возможность загружать растущее количество моделей. Очевидно, что это вертикальное масштабирование становится более проблематичным по мере увеличения требований к размеру базовой виртуальной машины.
Инициализация службы становилась еще одним узким местом, так как загрузка всех этих моделей при запуске занимала слишком много времени.

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

Преимущество описанного выше подхода в том, что:

Он может легко развернуть новый документ, не касаясь существующих документов, что означает более быстрое развертывание / загрузку.
Он быстрее масштабирует (увеличивает количество модулей) любую отдельную услугу. Если спрос на определенные документы будет расти, будет увеличиваться только эта индивидуальная услуга.
Мы не столкнемся с проблемой наличия большого узла с огромным ресурсом (ЦП / ОЗУ).

Мы также изменили стратегию и избежали запекания моделей в базовое изображение. Это значительно ускорило все развертывание из-за меньшего размера образа, который был загружен / извлечен в ECR во время развертывания.

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

Последний кусок головоломки решал проблему маршрутизации пользовательских запросов в соответствующую службу. Мы использовали Traefik в нашей инфраструктуре для обратного прокси, и он имеет возможность поддерживать сопоставление PathPrefix для маршрутизации. Поэтому мы определили правило Traefik для каждой службы на основе типа документа в пути, то есть пользователи должны передавать тип документа как часть URL-адреса, и Traefik перенаправит его на соответствующую службу. например, https://our-service.xyz/UK_PASSPORT/extract

Ограничения

Большинство ограничений, с которыми мы столкнулись, были связаны с AWS Batch:

AWS Batch всегда просматривает жестко запрограммированный тег изображения для создания нового контейнера (задания). Это ограничит возможность одновременной отправки нескольких заданий.
Консоль AWS Batch невыразительна, например, задание переходит в рабочее состояние, но вы не знаете, почему, и он не сообщает вам
Трудно настроить узел (особенно в старой версии GPU AMI), но после того, как он будет установлен, это будет менее болезненно.
Меньшая видимость / прямой доступ к контейнеру (выполнение задания), поскольку задание выполняется как контейнер на узле, управляемом ECS. В ECS не так много прозрачности по сравнению с облачными решениями, такими как Kubernetes.
Концепция минимального / максимального / желаемого ЦП сбивает с толку в пакетном режиме. Это добавляет сложности, когда требуемый виртуальный ЦП изменяется динамически, как только начинается новое задание.
Журналы контейнера по умолчанию перенаправляются в Cloudwatch. Однако журнал Cloudwatch нелегко прочитать. Вам нужно настроить лучшую службу ведения журнала на стороне.

Что дальше?

Поскольку наши модели обучаются с использованием TensorFlow, мы, естественно, ищем лучший вариант для обслуживания запросов. До сих пор мы видели некоторые многообещающие результаты от TensorFlow Serving, который снижает совокупное использование ресурсов моделями по сравнению с текущими моделями развертывания и быстрее обслуживает запросы.
Мы также узнали, насколько полезна такая инфраструктура и как с ее помощью можно масштабировать эксперименты с машинным обучением. Поэтому мы работаем над созданием инфраструктуры машинного обучения для наших исследователей, одновременно извлекая уроки, полученные с помощью AWS Batch. Следите за моей следующей статьей об инфраструктуре машинного обучения на Onfido.

Хотите помочь нам решить эти проблемы? Мы нанимаем инженеров-программистов, ученых-исследователей и других в Лондоне 🇬🇧 и Лиссабоне 🇵🇹