TL;DR В этой статье я объясню, как можно развернуть модель прямо в браузере из pytorch с помощью Onnjx. Эта работа была сделана год назад примерно за две недели.

https://narsil.github.io/2020/07/29/deploying-model-in-the-browser.html

Итак, когда мы демонстрируем глубокое обучение, обычно это подразумевает запуск моделей где-то в облаке. Иногда запуск этих моделей сам по себе обходится довольно дорого. Обучение GPT-3 стоило около 10 миллионов, но представьте, сколько будет стоить запуск, если бы он был доступен для широкой публики!

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

Задний план

Около года назад я работал в Набла (с тех пор они развернулись). Мы смотрели, насколько эффективна 3D-оценка позы лица. Это означает, что модели обнаруживают лица с глубиной, которая не была такой повсеместной, как обычное обнаружение лиц в 2D.

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

Начнем с 3DDFA

Итак, 3DDFA — это улучшенная реализация Pytorch этой статьи. Мы остановились на этой реализации, потому что она была лучшей из доступных на тот момент.

Как это работает ?

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

Для лиц 3ddfa использует смесь BFM и Facewarehouse, для полного человеческого тела часто используется SMPL.

В случае лица все 3D-сканы регистрируют одни и те же 53490 вершин на разных лицах разных участников (100 мужчин и 100 женщин для BFM). Например, центральная вершина носа:

Затем у нас есть те же участники, с разными выражениями:

В итоге мы получаем 150 участников x 20 поз x 53490 вершин (3 реалов). Затем мы можем использовать PCA для уменьшения размерности в 2 ортогональных пространства, одно для формы, а другое для выражения лица. Так что любое лицо может быть выражено как.

На самом деле нам не хватает глобальной ориентации лица, которую можно представить как матрицу вращения, коэффициент масштабирования и смещение, которые мы просто выразим в виде одной матрицы R и вектора p.

Модель 3DDFA фактически будет предсказывать по фото R, S_params, W_params и p. Для справки: R — 9 чисел с плавающей запятой, S_params — 40 чисел с плавающей запятой, W_params — 10 чисел с плавающей запятой, а p — 3. поплавки. Таким образом, фактическая модель 3DDFA берет изображение размером 120x120 пикселей и возвращает вектор размером 62, представляющий лицо.

Фактическая архитектура 3ddfa на самом деле представляет собой простую мобильную сеть.

Давайте портируем модель в браузер.

Итак, давайте изолируем первые несколько строк в main.py, которые загружают модель.

import torch
import mobilenet_v1
# 1. load pre-tained model                                                    
checkpoint_fp = 'models/phase1_wpdc_vdc.pth.tar'                              
arch = 'mobilenet_1'                                                          
                                                                              
checkpoint = torch.load(checkpoint_fp, map_location=lambda storage, loc: storage)['state_dict']
model = getattr(mobilenet_v1, arch)(num_classes=62)  # 62 = 12(pose) + 40(shape) +10(expression)
                                                                              
model_dict = model.state_dict()                                               
# because the model is trained by multiple gpus, prefix module should be removed
for k in checkpoint.keys():                                                   
    model_dict[k.replace('module.', '')] = checkpoint[k]                      
model.load_state_dict(model_dict)

Теперь добавим несколько строк для экспорта модели в Onnx

# Batch size, C, H, W
dummy_input = torch.zeros((1, 3, 120, 120) )

torch.onnx.export(model, dummy_input, "3ddfa.onnx", verbose=True, input_names=["input"], output_names=["params"])

Хорошо, теперь у нас есть 3ddfa.onnx файловая модель в нашем каталоге.

Давайте попробуем запустить его в браузере, следуя OnnxJS Getting Started, и напишем index.html файл:

<html>
  <head> </head>

  <body>
    <!-- Load ONNX.js -->
    <script src="https://cdn.jsdelivr.net/npm/onnxjs/dist/onnx.min.js"></script>
    <!-- Code that consume ONNX.js -->
    <script>
      // create a session
      const myOnnxSession = new onnx.InferenceSession();
      // load the ONNX model file

      function getInputs(){
          const x = new Float32Array(1 * 3 * 120 * 120).fill(1);
          const tensorX = new onnx.Tensor(x, 'float32', [1, 3, 120, 120]);
          return [tensorX];
      }

      myOnnxSession.loadModel("./3ddfa.onnx").then(() => {
        // generate model input
        const inferenceInputs = getInputs();
        // execute the model
        myOnnxSession.run(inferenceInputs).then((output) => {
          // consume the output
          const outputTensor = output.values().next().value;
          console.log(`model output tensor: ${outputTensor}.`);
        });
      });
    </script>
  </body>
</html>

Теперь давайте запустим локальный сервер с python -m http.server и перейдем к http://localhost:8000, чтобы увидеть вашу консоль. Упс! Мы получили ошибку: Uncaught (in promise) TypeError: cannot resolve operator 'Shape' with opsets: ai.onnx v9.

На самом деле onnx.js не может выводить динамические формы так же, как PyTorch. Мы можем отредактировать это, жестко закодировав некоторые значения в Reshape (что поддерживается).

В mobilenet_v1.py строке 144 измените x = x.view(x.shape(0), -1) на x.view(1, 1024). Теперь повторно экспортируйте файл onnx.

Это работает ! (Ну, у нас есть ценности.)

Внедрение этой модели в настоящий демонстрационный продукт.

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

Получить данные веб-камеры

var that = this
if (navigator.mediaDevices.getUserMedia) {
  navigator.mediaDevices.getUserMedia({video: {facingMode: 'user'}})
    .then(function(stream) {
            that.video.srcObject = stream;
            that.video.play().then(() => {
                    // That draws the video on a canvas.
                    that.facedetector.loadModel()
                    that.loop()
            }).catch((e) => {
                alert("Error launching webcam " + e)
            });

    }).catch(function(e){
        alert("No webcam detected " + e);
    });
}

Получить 3d лица из изображения

Мы используем детектор лиц, чтобы определить, где находятся лица на текущем изображении, а затем запускаем нашу модель для каждого лица на изображении.

var ctx = outcanvas.getContext('2d')!;
const detections = await faceapi.detectAllFaces(
  incanvas,
  new faceapi.TinyFaceDetectorOptions(),
);

var vertices = [];
for (i = 0; i < detections.length; i += 1) {
  var detection = detections[i];
  // We get a 1x3x120x120 Tensor here.
  // We could batch that in theory, but simplicity here.
  const inferenceInputs = this.getInputs(incanvas, detection);
  const outputData = await this.session!.run([inferenceInputs]);
  const output = outputData.values().next().value;

  // We need a reconstruction.
  const face_vertices = this.reconstruct68(output);
  
  // Ellipsed code where we fuse various meshes to only run a single rotation
  // and render process (we need to render to occlude the glasses in 3d with
  // a transparent mesh of the face.
}

ctx.drawImage(incanvas, 0, 0);
// Drawing back the occluded glasses on the webcam canvas.
this.drawGlasses(this.scene, outcanvas);

Восстановите вершины.

Как оказалось, поворот (R) + смещение (p), которые мы видели в первой части, на самом деле не то, чем кажется, R это просто матрица 3x3, ничто не заставляет ее вращаться (имеется в виду определитель 1). Как оказалось, модель довольно сильно меняет определитель этой матрицы с течением времени. Использование его в нашей демонстрации означало бы, что очки будут то увеличиваться, то уменьшаться и все время деформироваться.

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

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

Итак, что мы собираемся сделать, это реконструировать 68 вершин из граней модели. (нам нужно больше, чем 3, чтобы стабилизировать решение, 53 КБ слишком много) и solve уравнение

Где R имеет определитель 1, s скаляр и p вектор

Это 13 параметров (9 + 1 + 3) для 68 уравнений. Насколько я знаю, инвертирование этой системы не очень практично в javasript, поэтому мы собираемся просто запустить градиентный спуск. Это должно быть быстрее, чем инверсия при последующих обновлениях (когда решение близко к предыдущему решению).

Полный код выглядит немного пугающе (определенно не так, как я бы решил это сейчас), но в основном сводится к ручной работе с шагами градиента.

Добавьте очки на эту модель.

Если бы мы сразу добавили сетку очков, у нас возникла бы проблема с окклюзией, когда задняя ветвь была бы видна поверх вашего лица. Что мы собираемся сделать, так это создать сгруппированную сетку из очков и прозрачного лица (мы будем использовать среднее нормальное лицо, чтобы уменьшить объем вычислений), чтобы 3D-рендеринг закрывал заднюю часть. филиал очков.

Добавляем общее лицо:

public add_face(transparent?: boolean) {
  if (transparent === undefined) {
    transparent = true;
  }
  const self = this;
    this.scene.load(process.env.PUBLIC_URL + '/3dmodels/face.fbx').then(face => {
    if (transparent) {
      const mesh = face.children[0] as THREE.Mesh;
      mesh.renderOrder = -1;
      const material = mesh.material as THREE.Material;
  
      // Makes the face occluding, but we write the background
      // Color, so alpha instead of texture.
      material.colorWrite = false;
    }
    self.add(face);
    self.reset_clones();
  });
}

Добавление очков

public addElement(element: string) {
  this.group.add_element(process.env.PUBLIC_URL + `/3dmodels/${element}.fbx`);
}

Весь конвейер немного сложен для повторного использования мешей и эффективного запуска (все еще довольно беспорядочный код), но в этом суть, добавьте 2 меша в одну и ту же сцену и нажмите кнопку рендеринга. Фактически мы добавляем все лица и все очки в одну и ту же глобальную сцену и рендерим только один раз.

Вывод

Итак, мы увидели, как мы поместили модель из кода рабочего документа в реальный продукт. Необходимо помнить, что важно понимать, как бумага на самом деле работает, а не делает вид, что работает, здесь было очень важно понимать R и p и s влияет на нашу реконструкцию. Это также демонстрирует, что авторы статьи не осознавали, что то, что они считали позой, на самом деле содержит много искажений, поэтому это не работает так, как они ожидали.

Это также показывает, что внедрение машинного обучения в продукт даже для очень простой демонстрации требует понимания многих других технологий. Здесь это означало получение веб-камеры, упрощение слоев Onnx, понимание Three.js для освоения окклюзии и т. д. Но, перейдя в полноценный браузер, вы можете запустить демо и никогда не беспокоиться о стоимости!

Проверьте Полный исходный код для всех причуд.

Первоначально опубликовано наhttps://narsil.github.io/2020/07/29/deploying-model-in-the-browser.html