Scala + MXNet = Микросервис с нейронкой в проде

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

Основной, на мой взгляд, недостаток всех этих примеров — ограниченность возможностей. Вы взяли пример, — пусть даже с базовой нейронной сетью, которую предлагает автор, — запустили его, возможно, он даже заработал, а что дальше? Как сделать так, чтобы этот незамысловатый код начал работать на production-сервере? Как его обновлять и поддерживать? Вот тут и начинается самое интересное. Мне не удалось найти полного описания процесса от момента «ну вот, ML-инженер обучил нейронную сеть» до «наконец-то мы выкатили это в production». И я решил закрыть этот пробел.

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

Сразу отвечаю на вопрос «А почему не на Python?»: мы используем Scala для production-решений из-за более удобного и стабильного написания многопоточного кода.

Содержание

1. Постановка задачи
2. Используемые технологии
3. Подготовка базового Docker-контейнера
4. Структура проекта
5. Загрузка нейронной сети
6. Реализация REST API
7. Тестирование
8. Сборка микросервиса на основе базового образа
9. Запуск микросервиса на production-сервере с GPU
Заключение
Ссылки

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

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

2. Используемые технологии

• Scala 2.12.7 + минимальный набор дополнительных библиотек, Sbt 1.2.6 с плагином Sbt-pack 0.12 для сборки исходников.
• MXNet 1.3.1 (последняя стабильная версия на момент написания статьи), собранная для Scala 2.12.
• Сервер с видеокартами Nvidia.
• Cuda 9.0 и Cudnn 7, установленные на сервере.
• Java 8 для запуска скомпилированного кода.
• Docker для удобства сборки, доставки и запуска микросервиса на сервере.

3. Подготовка базового Docker-контейнера

Для нашего микросервиса потребуется базовый Docker-образ, в котором будет установлено минимальное количество необходимых для запуска зависимостей. Для сборки воспользуемся образом с дополнительно установленной Sbt. Да, сборку самих исходников будем выполнять не в локальном окружение, а в Docker-контейнере. Это облегчит в дальнейшем переход на сборку через CI, например через gitlab CI.

Структура папок:

| ----- install
|   | ----- java8.sh
|   | ----- mxnet_2_12.sh
|   | ----- opencv.sh
|   | ----- sbt.sh
|   | ----- scala.sh
|   | ----- timeZone.sh
| ----- scala-mxnet-cuda-cudnn
| ----- Dockerfile.2.12-1.3.1-9-7-builder
| ----- Dockerfile.2.12-1.3.1-9-7-runtime

Dockerfile.2.12-1.3.1-9-7-runtime

Этот образ будет использоваться для финального запуска микросервиса. В его основе лежит официальный образ от Nvidia с предустановленными CUDA 9.0 и CUDNN 7. В документации к MXNet 1.3.1 заявлена работа с CUDA 8.0, но, как показала практика, с версией 9.0 тоже всё прекрасно работает, и даже чуть быстрее.

Дополнительно мы установим в этот образ Java 8, MXNet 1.3.1 (будем собирать его под Scala 2.12), OpenCV 3.4.3 и Linux-утилиту для задания часового пояса.

# Собираем образ на основе базового образа от Nvidia с cuda 9.0 и cudnn 7
FROM nvidia/cuda:9.0-cudnn7-devel AS builder

# Объявление переменных окружения
ENV MXNET_VERSION 1.3.1
ENV MXNET_BUILD_OPT "USE_OPENCV=1 USE_BLAS=openblas USE_CUDA=1 USE_CUDA_PATH=/usr/local/cuda USE_CUDNN=1"
ENV CUDA_STUBS_DIR "/usr/local/cuda-9.0/targets/x86_64-linux/lib/stubs"
ENV OPEN_CV_VERSION 3.4.3
ENV OPEN_CV_INSTALL_PREFIX /usr/local
ENV JAVA_HOME /usr/lib/jvm/java-8-oracle/
ENV TIME_ZONE Europe/Moscow

# Копирование скриптов для установки
COPY install /install
RUN chmod +x -R /install/*

#  Запуск установки
RUN apt-get update
WORKDIR /install
RUN ./timeZone.sh ${TIME_ZONE}
RUN ./java8.sh
RUN ./mxnet_2_12.sh ${MXNET_VERSION} "${MXNET_BUILD_OPT}" ${CUDA_STUBS_DIR}
RUN ./opencv.sh ${OPEN_CV_VERSION} ${OPEN_CV_INSTALL_PREFIX}

# Удаление мусора из образа
RUN apt-get autoclean -y && 
rm -rf /var/cache/* /install

# Сборка финального образа и перенос данных
FROM nvidia/cuda:9.0-cudnn7-devel
COPY --from=builder --chown=root:root / /

Скрипты timeZone.sh java8.sh и opencv.sh достаточно тривиальные, поэтому не буду останавливаться на них подробно, они представлены ниже.

timeZone.sh

#!/bin/sh
# Получение название временной зоны из первого передаваемого параметра
TIME_ZONE=${1}
# Установка утилиты и задание временной зоны
apt-get install -y tzdata && 
   ln -sf /usr/share/zoneinfo/$TIME_ZONE /etc/localtime && 
   dpkg-reconfigure -f noninteractive tzdata

java8.sh

#!/bin/sh

# Установка Java 8
apt-get install -y software-properties-common && 
   add-apt-repository ppa:webupd8team/java -y && 
   apt-get update && 
   echo "oracle-java8-installer shared/accepted-oracle-license-v1-1 select true" | debconf-set-selections && 
   apt-get install -y oracle-java8-installer

opencv.sh

#!/bin/sh

# Получение версии OpenCV из первого передаваемого параметра
OPEN_CV_VERSION=${1}

# Получение префикса установки из второго параметра скрипта
OPEN_CV_INSTALL_PREFIX=${2}
OPEN_CV_TAR="http://github.com/opencv/opencv/archive/${OPEN_CV_VERSION}.tar.gz"

# Установка OpenCV
apt-get install -y wget build-essential cmake && 
   wget -qO- ${OPEN_CV_TAR} | tar xzv -C /tmp && 
   mkdir /tmp/opencv-${OPEN_CV_VERSION}/build && 
   cd /tmp/opencv-${OPEN_CV_VERSION}/build && 
   cmake -DBUILD_JAVA=ON -DCMAKE_INSTALL_PREFIX:PATH=${OPEN_CV_INSTALL_PREFIX} .. && 
   make -j$((`nproc`+1)) && 
   make install && 
   rm -rf /tmp/opencv-${OPEN_CV_VERSION}

С установкой MXNet не всё так просто. Дело в том, что все сборки этой библиотеки для Scala делаются на базе компилятора версии 2.11, и это обоснованно, так как в состав библиотеки входит модуль для работы со Spark, который, в свою очередь, написан на Scala 2.11. Учитывая то, что мы в разработке используем Scala 2.12.7, собранные библиотеки нам не подходят, а спуститься на версию 2.11.* мы не можем ввиду большого количества уже написанного кода на новой версии Scala. Что делать? Получать массу удовольствия, собирая MXNet из исходников под нашу версию Scala. Ниже я приведу скрипт для сборки и установки MXNet 1.3.1 для Scala 2.12.* и прокомментирую основные моменты.

mxnet_2_12.sh

#!/bin/sh

# Получение версии MXNet из первого передаваемого параметра
MXNET_VERSION=${1}

# Получение опций для сборки С++ библиотек MXNet из второго передаваемого параметра
MXNET_BUILD_OPT=${2}

# Получение пути до директории с библиотеками CUDA из третьего передаваемого параметра
CUDA_STUBS_DIR=${3}
LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${CUDA_STUBS_DIR}"

# Установка зависимостей для компиляции и сборки MXNet и установка
apt-get install -y git build-essential libopenblas-dev libopencv-dev maven cmake && 
   git clone -b ${MXNET_VERSION} --recursive https://github.com/dmlc/mxnet /tmp/mxnet && 
   cd /tmp/mxnet && 
   make -j $(nproc) ${MXNET_BUILD_OPT} && 
   ln -s ${CUDA_STUBS_DIR}/libcuda.so ${CUDA_STUBS_DIR}/libcuda.so.1 && 
   sed -rim 's/([a-zA-Z])_2.11/1_2.12/g' $(find scala-package -name pom.xml) && 
   sed -im 's/SCALA_VERSION_PROFILE := scala-2.11/SCALA_VERSION_PROFILE := scala-2.12/g' Makefile && 
   sed -im 's/<module>spark</module>/<!--<module>spark</module>-->/g' scala-package/pom.xml && 
   make scalapkg ${MXNET_BUILD_OPT} && 
   mkdir -p /usr/local/share/mxnet/scala/linux-x86_64-gpu && 
   mv /tmp/mxnet/scala-package/assembly/linux-x86_64-gpu/target/mxnet-full_2.12-linux-x86_64-gpu-${MXNET_VERSION}-SNAPSHOT.jar /usr/local/share/mxnet/scala/linux-x86_64-gpu/mxnet-full_2.12-linux-x86_64-gpu-${MXNET_VERSION}-SNAPSHOT.jar && 
   rm -rf /tmp/mxnet && rm -rf /root/.m2

Самое интересное начинается с этой строки:

ln -s ${CUDA_STUBS_DIR}/libcuda.so ${CUDA_STUBS_DIR}/libcuda.so.1 && 

Если запустить сборку MXNet как в инструкции, то мы получим ошибку. Компилятор не может найти библиотеку libcuda.so.1, поэтому сделаем ссылку с библиотеки libcuda.so на libcuda.so.1. Пусть это вас не смущает, при запуске на production-сервере мы будем подменять эту библиотеку локальной. Также обратите внимание, что в LD_LIBRARY_PATH добавлен путь к CUDA-библиотекам из переменной окружения CUDA_STUBS_DIR. Если этого не сделать, то сборка тоже завалится.

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

sed -rim 's/([a-zA-Z])_2.11/1_2.12/g' $(find scala-package -name pom.xml) && 
    sed -im 's/SCALA_VERSION_PROFILE := scala-2.11/SCALA_VERSION_PROFILE := scala-2.12/g' Makefile && 
    sed -im 's/<module>spark</module>/<!--<module>spark</module>-->/g' scala-package/pom.xml && 
    make scalapkg ${MXNET_BUILD_OPT} && 

И тут же комментируется зависимость от модуля работы со Spark. Если этого не сделать, библиотека не соберётся.

Далее запускаем сборку, как указано в инструкции, копируем собранную библиотеку в общую папку и удаляем всякий мусор, который в процессе сборки накачал Maven (если этого не сделать, то финальный образ подрастёт в размере примерно на 3-4 Гб, что может заставить ваших DevOps-ов нервничать).

Собираем образ, находясь в корневой директории проекта (см. Структура папок):

your@pc$ docker build -f Dockerfile.2.12-1.3.1-9-7-runtime -t entony/scala-mxnet-cuda-cudnn:2.12-1.3.1-9-7-runtime .

Напомню на всякий случай, что точка в конце говорит о том, что мы делаем сборку в рамках контекста текущей директории.

Теперь пришло время поговорить об образе для сборки.

Dockerfile.2.12-1.3.1-9-7-builder

# Тут уже в качестве базового образа будет использоваться runtime-образ, который мы собрали
FROM entony/scala-mxnet-cuda-cudnn:2.12-1.3.1-9-7-runtime

# Объявляем переменные окружения
ENV SCALA_VERSION 2.12.7
ENV SBT_VERSION 1.2.6

# Копирование скриптов для установки
COPY install /install
RUN chmod +x -R /install/*

# Устанавливаем инструменты для сборки и зависимости
RUN apt-get update && 
   cd /install && 
   ./scala.sh ${SCALA_VERSION} && 
   ./sbt.sh ${SBT_VERSION}

# Удаляем мусор
RUN rm -rf /install

Тут всё просто, для запуска нашего микросервиса нам не нужны Scala и Sbt, поэтому нет смысла тащить их в базовый образ для запуска. Поэтому создадим отдельный образ, который будет использоваться только для сборки. Скрипты scala.sh и sbt.sh достаточно тривиальные и подробно останавливаться на них я не буду.

scala.sh

#!/bin/sh

# Получение версии Scala из первого передаваемого параметра
SCALA_VERSION=${1}
SCALA_DEB="http://www.scala-lang.org/files/archive/scala-${SCALA_VERSION}.deb"

# Установка Scala
apt-get install -y wget && 
   wget -q ${SCALA_DEB} -O /tmp/scala.deb && dpkg -i /tmp/scala.deb && 
   scala -version && 
   rm /tmp/scala.deb

sbt.sh

#!/bin/sh

# Получение версии Sbt из первого передаваемого параметра
SBT_VERSION=${1}
SBT_DEB="http://dl.bintray.com/sbt/debian/sbt-${SBT_VERSION}.deb"

# Установка Sbt
apt-get install -y wget && 
   wget -q ${SBT_DEB} -O /tmp/sbt.deb && dpkg -i /tmp/sbt.deb && 
   sbt sbtVersion && 
   rm /tmp/sbt.deb

Собираем образ, находясь в корневой директории проекта (см. Структура папок):

your@pc$ docker build -f Dockerfile.2.12-1.3.1-9-7-builder -t entony/scala-mxnet-cuda-cudnn:2.12-1.3.1-9-7-builder .

В конце статьи есть ссылки на репозиторий со всеми этими файлами.

4. Структура проекта

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

Проект нашего микросервиса будет иметь следующую структуру:

| ----- dependencies
|   | ----- mxnet-full_2.12-linux-x86_64-gpu-1.3.1-SNAPSHOT.jar
| ----- models
|   | ----- resnet50_ssd_model-0000.params
|   | ----- resnet50_ssd_model-symbol.json
|   | ----- synset.txt
| ----- project
|   | ----- build.properties
|   | ----- plugins.sbt
| ----- src
|   | ----- main
|   |   | ----- resources
|   |   |   | ----- cat_and_dog.jpg
|   |   | ----- scala
|   |   |   | ----- simple.predictor
|   |   |   | ----- Config
|   |   |   | ----- Model
|   |   |   | ----- Server
|   |   |   | ----- Run
|   | ----- test
|   |   | ----- scala
|   |   |   | ----- simple.predictor
|   |   |   | ----- ServerTest
| ----- build.sbt
| ----- Dockerfile

Это стандартная структура Scala-проекта, за исключением директорий dependencies и models.
В директории dependencies находится библиотека MXNet для Scala. Её можно получить двумя способами:

• собрать MXNet на машине, на которой будете производить разработку (обратите внимание, что библиотека не кроссплатформенная, если соберёте её на Linux, она не будет работать на Mac OS),
• или вытащить её из Docker-образа, который мы собирали ранее. Если надумаете собирать MXNet в локальном окружение, то скрипт mxnet_2.12.sh вам в помощь.

Вытащить библиотек из Docker-образа можно вот так:

# Создаём директорию
your@pc$ mkdir dependencies

# Запускаем Docker-контейнер с перемонтированной директорией
your@pc$ docker run -it --rm -v $(pwd)/dependencies:/tmp/dependencies entony/scala-mxnet-cuda-cudnn:2.12-1.3.1-9-7-runtime

# Из контейнера копируем файл в перемонтированную директорию и выходим
ab38e73d93@root$ cp /usr/local/share/mxnet/scala/linux-x86_64-gpu/mxnet-full_2.12-linux-x86_64-gpu-1.3.1-SNAPSHOT.jar /tmp/dependencies/mxnet-full_2.12-linux-x86_64-gpu-1.3.1-SNAPSHOT.jar
ab38e73d93@root$ exit

# Вот оно, счастье!
your@pc$ ls dependencies/
mxnet-full_2.12-linux-x86_64-gpu-1.3.1-SNAPSHOT.jar

В директории models находятся файлы обученной нейронной сети, скачать их можно свободно следующим образом:

# Создаём директорию
your@pc$ mkdir models

# Скачиваем файлы в директорию
your@pc$ wget https://s3.amazonaws.com/model-server/models/resnet50_ssd/resnet50_ssd_model-symbol.json -P models
your@pc$ wget https://s3.amazonaws.com/model-server/models/resnet50_ssd/resnet50_ssd_model-0000.params -P models
your@pc$ wget https://s3.amazonaws.com/model-server/models/resnet50_ssd/synset.txt -P models

Далее кратко о файлах, которые не представляют особого интереса, но играют роль в проекте.

project/build.properties

# Просто версия Sbt, ничего больше
sbt.version = 1.2.6

project/plugins.sbt

# Зависимость от плагина sbt-pack
addSbtPlugin("org.xerial.sbt" % "sbt-pack" % "0.12")

src/main/resources/cat_and_dog.jpg

Такая вот замечательная картинка, на которой наша нейронная сеть будет искать котика и собачку.

build.sbt

enablePlugins(PackPlugin)

name := "simple-predictor"

version := "0.1"

scalaVersion := "2.12.7"

unmanagedBase := baseDirectory.value / "dependencies"

// Зависимости (минимальное количество для примера)
libraryDependencies ++= Seq(
  "org.json4s" %% "json4s-native" % "3.6.1",
  "org.scalatest" %% "scalatest" % "3.0.5" % Test,
  "org.scalaj" %% "scalaj-http" % "2.4.1" % Test
)

// Псевдоним для запуска микросервиса после сборки
packMain := Map("simple-predictor" -> "simple.predictor.Runs")

// Нам не нужен bat-файл, так как запуск будет внутри контейнера, а там Linux
packGenerateWindowsBatFile := false

// Ключи для запуска JVM
packJvmOpts := Map("simple-predictor" -> Seq( "-Xms3g", "-Xmx5g"))

simple.predictor.Config

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

package simple.predictor

import org.apache.mxnet.Context

import scala.util.Try

object Config {

  // Хост для запуска REST API
  val host: String = env("REST_HOST") getOrElse "0.0.0.0"

  // Порт для запуска REST API
  val port: Int = env("REST_PORT") flatMap (p => Try(p.toInt).toOption) getOrElse 8080
  // URL, на который будет отправляться POST-запрос с картинкой
  val entryPoint: String = env("REST_ENTRY_POINT") getOrElse "/predict"

  // Порог вероятности, после которого найденный объект считается ошибочным
  val threshold: Float = env("PROBABILITY_MORE") flatMap (p => Try(p.toFloat).toOption) getOrElse 0.5f

  // Префикс файлов нейронной сети относительно директории проекта
  val modelPrefix: String = env("MODEL_PREFIX") getOrElse "models/resnet50_ssd_model"

  // Эпоха нейронной сети (можно определить из названия ...-0000.params)
  val modemEpoch: Int = env("MODEL_EPOCH") flatMap (p => Try(p.toInt).toOption) getOrElse 0

  // Размер стороны картинки, который использовался для обучения сети, у нашей сети 512
  val modemEdge: Int = env("MODEL_EDGE") flatMap (p => Try(p.toInt).toOption) getOrElse 512

  // Контекст запуска, по умолчанию CPU (для разработки). В production будет GPU
  val context: Context = env("MODEL_CONTEXT_GPU") flatMap {
    isGpu => Try(if (isGpu.toBoolean) Context.gpu() else Context.cpu()).toOption
  } getOrElse Context.cpu()

private def env(name: String) = Option(System.getenv(name))

}

simple.predictor.Run

Объект Run — это точка входа в приложение.

package simple.predictor

import java.net.InetSocketAddress

// Импортируем объект с конфигурацией
import simple.predictor.Config._

object Run extends App {

  // Инициализация нейронной сети и REST-сервера
  val model = new Model(modelPrefix, modemEpoch, modemEdge, threshold, context)
  val server = new Server(new InetSocketAddress(host, port), entryPoint, model)

  // При нажатии Ctrl + C приложение будет завершаться
  Runtime.getRuntime.addShutdownHook(new Thread(() => server.stop()))

  // Запускаем сервер и отлавливаем исключения
  try server.start() catch {
    case ex: Exception => ex.printStackTrace()
  }

}

5. Загрузка нейронной сети

Нейронная сеть загружается в конструкторе класса simple.predictor.Model.

simple.predictor.Model

package simple.predictor

import java.awt.image.BufferedImage
import org.apache.mxnet._
import org.apache.mxnet.infer.ObjectDetector
import simple.predictor.Model.Prediction

class Model(prefix: String, epoch: Int, imageEdge: Int, threshold: Float, context: Context) {

  // Подготовка описания данных для инициализации модели
  val initShape = Shape(1, 3, imageEdge, imageEdge)
  val initData = DataDesc(name = "data", initShape, DType.Float32, Layout.NCHW)

  // Загрузка модели из файлов с заданным префиксом с указанным контекстом
  val model = new ObjectDetector(prefix, IndexedSeq(initData), context, Option(epoch))

  // Эта функция преобразует возвращаемые моделью данные в результирующий объект, который в дальнейшем будет преобразовываться в JSON
  private def toPrediction(originWidth: Int, originHeight: Int)(predict: (String, Array[Float])): Prediction = {
    val (objectClass, Array(probability, kx, ky, kw, kh)) = predict
    // Пересчёт координат с учётом реальных размеров изображения
    val x = (originWidth * kx).toInt
    val y = (originHeight * ky).toInt
    val w = (originWidth * kw).toInt
    val h = (originHeight * kh).toInt
    val width = if ((x + w) < originWidth) w else originWidth - x
    val height = if (y + h < originHeight) h else originHeight - y
    Prediction(objectClass, probability, x, y, width, height)
}

  // Прогоняет изображение через нейронную сеть, преобразует результат в объект и фильтрует результаты с вероятностью, меньше заданной в поле threshold
  def predict(image: BufferedImage): Seq[Prediction] =
    model.imageObjectDetect(image).head map toPrediction(image.getWidth, image.getHeight) filter (_.probability > threshold)

}

object Model {

  // Результирующий объект
  case class Prediction(objectClass: String, probability: Float, x: Int, y: Int, width: Int, height: Int)

}

В разделе Подготовка описания данных для инициализации модели вы говорите, что в нейронной сети будет работать с NDArray размерность 1 х 3 х 512 х 512, где 1 — это количество изображений, которые будут содержаться в NDArray, 3 — количество цветов, а 512 х 512 — размер изображения (значение imageEdge = 12 задано в объекте simple.predict.Config, это размер стороны изображения, который использовался при обучение нейронной сети). Всё это описание данных передаётся в объект ObjectDetector.

Ещё интересен раздел — Пересчёт координат с учётом реальных размеров изображения.

После прогона изображения через нейронную сеть результат имеет тип Seq[Seq[(String, Array[Float])]]. Первая коллекция содержит только один результат (формат данных определяется конкретной нейронной сетью), далее каждый элемент следующей коллекции представляет собой кортеж из двух элементов:

1. название класса («cat», «dog», …),
2. массив из пяти чисел с плавающей точкой: первое — вероятность, второе — коэффициент для вычисления координаты x, третье — коэффициент для вычисления координаты y, четвёртое — коэффициент для вычисления ширины прямоугольника, и пятое — коэффициент для вычисления высоты прямоугольника.

Для получения действительных значений координат и размеров прямоугольника нужно оригинальные ширину и высоту изображения умножить на соответствующие коэффициенты.
Позволю себе небольшое лирическое отступление на тему NDArray. Это многомерный массив, который MXNet создаёт в заданном контексте (CPU или GPU). При создании NDArray формируется С++-объект, операции с которым выполняются очень быстро (а если он создан в GPU-контексте, то практически мгновенно), но за такую скорость нужно платить. В результате (как минимум, в версии MXNet 1.3.1) нужно самостоятельно управлять памятью, выделяемой под NDArray, и не забывать выгружать из памяти эти объекты после завершения работы с ними. В противном случае будет значительная и достаточно быстрая утечка памяти, которую не очень удобно отслеживать, так как программы для профилирования JVM её не видят. Проблема с памятью усугубляется, если вы работаете в GPU-контексте, так как видеокарты не обладают большим объёмом памяти и приложение быстро вылетает по out of memory.

Как решать проблему с утечкой памяти?

В приведённом примере, в строке model.imageObjectDetect(image).head map toPrediction(image.getWidth, image.getHeight) filter (_.probability > threshold) для выполнения прогона изображения через нейронную сеть используется метод imageObjectDetect, который получает на вход BufferedImage. Все преобразования в NDArray и обратно выполняются внутри метода, и вам не нужно думать о проблемах высвобождения памяти. С другой стороны, перед преобразованием BufferedImage в NDArray выполняется ресайзинг под размер 512 х 512 и нормирование изображения методами объекта типа BufferedImage. Так получается чуть дольше, чем при использовании OpenCV, например, но зато решается проблема с высвобождением памяти после использования NDArray.

Можно, конечно использовать OpenCV и контролировать память самостоятельно, для этого нужно всего лишь вызывать у NDArray метод dispose, но в официальной документации MXNet для Scala это почему-то упомянуть забыли.

Также в MXNet есть не совсем удобный способ контроля утечки памяти, которая возникает из-за NDArray. Для этого нужно запустить приложение с JVM-параметром Dmxnet.traceLeakedObjects=true. Если MXNet заметит NDArray, который не используется, но висит в памяти, вы получите исключение с указанием, в какой строке кода находится злосчастный NDArray.

Мой совет: работайте с NDArray напрямую, внимательно следите за памятью и пишите нормализацию самостоятельно, предварительно уточнив, по какому алгоритму это делал ML-инженер при обучении нейронной сети, а то результаты будут совсем разные. У объекта ObjectDetector есть метод objectDetectWithNDArray, в который можно передавать NDArray. Для реализации более универсального подхода к загрузке нейронной сети рекомендую использовать объект org.apache.mxnet.module.Module. Ниже приведён пример использования.

import org.apache.mxnet._
import org.apache.mxnet.io.NDArrayIter

// Загрузка и инициализация нейронной сети
val model: Module = {
  val model = Module.loadCheckpoint(modelPrefix, modelEpoch, contexts = contexts)
  model.bind(
  forTraining = false,
  inputsNeedGrad = false,
  forceRebind = false,
  dataShape = DataDesc(name = "data", Shape(1, 3, 512, 512), DType.Float32, Layout.NCHW))
  model.initParams()
  model
}

// NDArray размерностью 1 х 3 х 512 х 512
val image: NDArray = ???

// Подготовка dataBatch для прогона через нейронную сеть
val iterator = new NDArrayIter(IndexedSeq(image))
val dataBatch = iterator.next()
image.dispose()

// Получение результата
val result: Seq[Array[Float]] = model.predict(dataBatch) map { ndArray =>
  val array = ndArray.toArray
  ndArray.dispose()
  array
}
dataBatch.dispose()

6. Реализация REST API

За реализацию REST API отвечает класс simple.predictor.Server. Сервер создан на основе входящего в Java HTTP-сервера.

simple.predictor.Server

package simple.predictor

import java.net.InetSocketAddress

import com.sun.net.httpserver.{HttpExchange, HttpServer}
import javax.imageio.ImageIO
import org.json4s.DefaultFormats
import org.json4s.native.Serialization

class Server(address: InetSocketAddress, entryPoint: String, model: Model) {

  // Будем использовать HTTP-сервер, который входит в поставку java
  private val server = HttpServer.create(address, 0)

  // Вешаем обработчик событий на наш URL
  server.createContext(entryPoint, (http: HttpExchange) => {

    // Проверяем заголовок HTTP-запроса на наличие нужного контента
    val header = http.getRequestHeaders
    val (httpCode, json) = if (header.containsKey("Content-Type") && header.getFirst("Content-Type") == "image/jpeg") {

      // В случае успеха выкачиваем картинку и прогоняем через нейронную сеть, возвращаем результат с кодом ответа 200
      val image = ImageIO.read(http.getRequestBody)
      val predictionSeq = model.predict(image)
      (200, Map("prediction" -> predictionSeq))
    } else (400, Map("error" -> "Invalid content")) // Иначе возвращаем ошибку с кодом ответа 400

    // Сериализуем результат в JSON и отправляем клиенту
    val responseJson = Serialization.write(json)(DefaultFormats)
    val httpOs = http.getResponseBody
    http.getResponseHeaders.set("Content-Type", "application/json")
    http.sendResponseHeaders(httpCode, responseJson.length)
    httpOs.write(responseJson.getBytes)
    httpOs.close()

  })

def start(): Unit = server.start()

def stop(): Unit = server.stop(0)

}

7. Тестирование

Для проверки запустим сервер и отправим тестовую картинку src/main/resources/cat_and_dog.jpg. Распарсим полученных JSON от сервера, проверим, сколько и каких объектов нейронная сеть нашла на картинке, и обведём объекты на картинке.

simple.predictor.ServerTest

package simple.predictor

import java.awt.{BasicStroke, Color, Font}
import java.awt.image.BufferedImage
import java.io.{ByteArrayOutputStream, File}
import java.net.InetSocketAddress

import javax.imageio.ImageIO
import org.scalatest.{FlatSpec, Matchers}
import scalaj.http.Http
import org.json4s.{DefaultFormats, Formats}
import org.json4s.native.JsonMethods.parse
import simple.predictor.Config._
import simple.predictor.Model.Prediction

import scala.concurrent.Future
import scala.concurrent.ExecutionContext.Implicits.global

class ServerTest extends FlatSpec with Matchers {

  implicit val formats: Formats = DefaultFormats

  "Service" should "find a cat and a dog on photo" in {

    // Инициализируем нейронную сеть и сервер
    val model = new Model(modelPrefix, modemEpoch, modemEdge, threshold, context)
    val server = new Server(new InetSocketAddress(host, port), entryPoint, model)

    // Запускаем сервер в отдельном потоке
    Future(server.start())
    Thread.sleep(5000)

    // Загружаем картинку и сериализуем её в массив байтов
    val image = ImageIO.read(getClass.getResourceAsStream("/cat_and_dog.jpg"))
    val byteOS = new ByteArrayOutputStream()
    ImageIO.write(image, "jpg", byteOS)
    val data = byteOS.toByteArray

    // Отправляем картинку на сервер и проверяем, что ответ имеет статус 200
    val response = Http(s"http://$host:$port$entryPoint").header("Content-Type", "image/jpeg").postData(data).asString
    response.code shouldEqual 200

    // Парсим JSON-ответ, проверяем, что он содержит два найденных объекта
    val prediction = parse(response.body) \ "prediction"
    prediction.children.size shouldEqual 2

    // Получаем список названий класса объектов, проверяем, что первый из них кот, а второй собака
    val objectClassList = (prediction \ "objectClass").children map (_.extract[String])
    objectClassList.head shouldEqual "cat"
    objectClassList.tail.head shouldEqual "dog"

    // Получаем координаты прямоугольников, выделяющих объекты
    val bBoxCoordinates = prediction.children.map(_.extract[Prediction])

    // Обводим найденные объекты, выводим названия и вероятности
    val imageWithBoundaryBoxes = new BufferedImage(image.getWidth, image.getHeight, image.getType)
    val graph = imageWithBoundaryBoxes.createGraphics()
    graph.drawImage(image, 0, 0, null)
    graph.setColor(Color.RED)
    graph.setStroke(new BasicStroke(5))
    graph.setFont(new Font(Font.SANS_SERIF, Font.TRUETYPE_FONT, 30))
    bBoxCoordinates foreach { case Prediction(obj, prob, x, y, width, height) =>
    graph.drawRect(x, y, width, height)
    graph.drawString(s"$obj, prob: $prob", x + 15, y + 30)
  }
  graph.dispose()

  // Сохраняем картинку с выделенными объектами в корень проекта
  ImageIO.write(imageWithBoundaryBoxes, "jpg", new File("./test.jpg"))

  }

}

После запуска тест проходит успешно, ниже представлена сохранённая картинка с выделенными объектами.

8. Сборка микросервиса на основе базового образа

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

Dockerfile

# Собираем на основе образа с установленной Sbt
FROM entony/scala-mxnet-cuda-cudnn:2.12-1.3.1-9-7-builder AS builder

# Создаём необходимые каталоги и копируем данные
RUN mkdir /tmp/source /tmp/source/dependencies
COPY project /tmp/source/project
COPY src /tmp/source/src
COPY build.sbt /tmp/source/build.sbt

# Делаем ссылку на библиотеку MXNet, которая внутри образа и запускаем сборку
RUN ln -s /usr/local/share/mxnet/scala/linux-x86_64-gpu/mxnet-full_2.12-linux-x86_64-gpu-1.3.1-SNAPSHOT.jar /tmp/source/dependencies/mxnet-full_2.12-linux-x86_64-gpu-1.3.1-SNAPSHOT.jar && 
cd /tmp/source/ && sbt pack

# Создаём чистый образ с микросервисом
FROM entony/scala-mxnet-cuda-cudnn:2.12-1.3.1-9-7-runtime

# Добавляем в LD путь к Cuda библиотекам и Java
ENV LD_LIBRARY_PATH /usr/local/cuda-9.0/targets/x86_64-linux/lib/stubs:/usr/local/share/OpenCV/java

# Наша модель теперь будет лежать не в корне проекта а в /opt/app/models
ENV MODEL_PREFIX "/opt/app/models/resnet50_ssd_model"

# Создаём директория с приложением и копируем туда модель и собранный микросервис
RUN mkdir -p /opt/app
COPY --from=builder --chown=root:root /tmp/source/target/pack /opt/app
COPY models /opt/app/models

# Запускаем микросервис про старте контейнера
ENTRYPOINT /opt/app/bin/simple-predictor

Собираем микросервис и загружаем его в докер репозиторий

# Сборка, запускаем в папку проекта, где находится Dockerfile
your@pc$ docker build -f Dockerfile -t entony/simple-predictor:1.0.0 .

# Загрузка в docker hub
your@pc$ docker push entony/simple-predictor:1.0.0

9. Запуск микросервиса на production-сервере с GPU

Предполагается, что образ микросервиса загружен в docker hub, на сервере есть одна видеокарта Nvidia, порт 8080 свободен и установлены Docker, Cuda 9.0 и Cudnn 7.

# Скачиваем образ микросервиса из Docker hub
your@server-with-gpu$ docker pull entony/simple-predictor:1.0.0

# Запускаем в режиме демона
your@server-with-gpu$ docker run -d 
    -p 8080:8080 
    -e MODEL_CONTEXT_GPU=true 
    -e MXNET_CUDNN_AUTOTUNE_DEFAULT=0 
    --name 'simple_predictor' 
    --device /dev/nvidia0:/dev/nvidia0 
    --device /dev/nvidiactl:/dev/nvidiactl 
    --device /dev/nvidia-uvm:/dev/nvidia-uvm 
    -v /usr/lib/x86_64-linux-gnu/libcuda.so.1:/usr/local/cuda-9.0/targets/x86_64-linux/lib/stubs/libcuda.so.1:ro 
    -v /usr/lib/nvidia-396/libnvidia-fatbinaryloader.so.396.54:/usr/local/cuda-9.0/targets/x86_64-linux/lib/stubs/libnvidia-fatbinaryloader.so.396.54:ro 
    entony/simple-predictor:1.0.0

Из интересного тут можно отметить пробрасывание видеокарт в Docker-контейнер через параметры --device и пробрасывание Cuda-библиотек через параметр -v.

Переменная окружения MODEL_CONTEXT_GPU определяет запуск модели в GPU-контексте, а MXNET_CUDNN_AUTOTUNE_DEFAULT отключает автоподбор параметров нейронной сети и попытки оптимизировать её (на мой взгляд, на скорость работы это не влияет, но раз в документации написано, пускай будет).

Отправим тестовый запрос с нашей картинкой:

# Запрос
your@server-with-gpu$ curl -X POST -H 'Content-Type: image/jpeg' --data-binary '@src/main/resources/cat_and_dog.jpg' http://0.0.0.0:8080/predict

# Ответ
{
  "prediction":[
    {
      "objectClass":"cat",
      "probability":0.9959417,
      "x":72,"y":439,
      "width":950,
      "height":987
    },
    {
      "objectClass":"dog",
      "probability":0.81277525,
      "x":966,
      "y":100,
      "width":870,
      "height":1326
    }
  ]
}

Заключение

MXNet показал себя неплохо, хотя пришлось поэкспериментировать из-за не очень качественной документации. Всё равно, общее впечатление осталось позитивное, нейросеть стоит того, чтобы тащить её в production.

На мой взгляд, удобно, что у MXNet есть поддержка нескольких языков, можно обучить нейронную сеть на Python и топом внедрить её в production на Scala, Java или С++.

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

Надеюсь, что прочтение этой статьи пошло вам на пользу. При поиске объектов на фото ни одно животное не пострадало. Спасибо за внимание.

Ссылки

• Git с базовыми Docker-образами
• Git с кодом из примера
• Описание похожего примера из документации MXNet

Автор: Anton Yurchenko

Источник