При работе с TensorFlowJS можно выделить три направления его использования (рисунок 1):

Использование обученных моделей без модификаций топологий и без их переобучений
Использование предварительно обученную модель с возможной модификацией ее топологии или переобучении на базе новой тренировочной выборки данных, которая полностью соответствует решаемой вашей задачи
Создание и обучение моделей с нуля

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

1. Использование моделей через абстрактный АПИ

Этот тот случай, когда вам не потребуются какие-то глубокие знания по машинному обучению в принципе. Вам даже не придется производить никаких манипуляций с входными данными для модели. Например, если искомая модель принимает на вход изображение размерностью 240x240 пикселей, а вы хотели бы использовать изображение 20x30, то вам бы пришлось делать манипуляции над изображением, чтобы она была совместима с моделью. Также достаточно часто, модели требуют нормализацию входных данных (значения в одном цветом канале для пикселей изменяются от 0 до 256, однако для лучшей сходимости модели могут иногда требовать, чтобы величина каждого пиксела была в интервале [0, 1] или [-1, 1]).

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

Рисунок 2 – Классификация моделей по типу решаемых задач

На момент написания статьи (октябрь 2020), Google представил 13 официальных моделей для использования в открытом доступе. Весь список моделей может быть найден тут. В списке вы можете найти модели, решающие следующие типы задач (рисунок 2):

классификация изображений: MobileNet - классификация изображений между 1000 категорий, при этом на изображении должен находится объект одного класса с нейтральным фоном (ссылка)
обнаружение объектов на изображении с указанием пространственных координат: Coco-SSD – определение объектов на изображении с указанием окна, в котором находится объект; модель может распознавать 80 разных категорий объектов (ссылка)
определение положения человеческого тела: BodyPix – пространственное определение частей тела (руки, ноги, плечи, глаза, уши), а также с помощью модели можно построить маску положения тела. На изображении одновременно может находиться несколько человек (ссылка)
распознавание голосовых команд: SpeechCommands – распознавание звуковых команд; в исходнос состоянии модель может распознавать команды из словаря из 20 слов на английском языке, например: 'go', 'stop', 'yes', 'no' (ссылка)
текстовая классификация: Toxicity – определяет содержит ли сообщения не приемлемый контент, содержащий оскорбления, не уважение, непристойные выражения с сексуальным содержанием (ссылка)
текстовая классификация: Toxicity – определяет содержит ли сообщения не приемлемый контент, содержащий оскорбления, не уважение, непристойные выражения с сексуальным содержанием (ссылка)
- .

Давайте рассмотрим АПИ конкретной модели Coco-SSD и каким образом оно может быть использовано:

import * as tf from "@tensorflow/tfjs";
import * as cocoSsd from "@tensorflow-models/coco-ssd";
// other imports


export default () => {
    const {videoRef} = useVideoStream();

    const canvasRef = useRef<HTMLCanvasElement>(null);
    const modelRef = useRef<ObjectDetection>();

    function detectFrame() {
        if (modelRef.current && canvasRef.current && videoRef.current) {
            modelRef.current.detect(videoRef.current)
                .then(objects => {
                    buildObjectReactangle(canvasRef.current, objects);
                    window.requestAnimationFrame(detectFrame);
                });
        }
    }

    useEffect(() => {
        cocoSsd.load().then(model => {
            modelRef.current = model;
            detectFrame();
        });
    }, []);

    return (
        <div style={{position: 'relative'}}>
            <video ref={videoRef} width={640} height={480}/>
            <canvas ref={canvasRef} width={640} height={480}/>
        </div>
    );
}

В первую очередь, необходимо сделать импорт модели и библиотеки @tensorflow/tfjs, этот тот единственный след, говорящий, что под капотом модель использует TensorFlowJS. При первом рендеринге React компонента необходимо загрузить модель: cocoSsd.load (строки 22-27). Как только модель будет загружена (это может занять некоторое время), можно запускать процесс обработки кадров, получаемые с видео потока вызовом detectFrame функции (строка 25). Для получения метаинформации о положении объектов на изображении – достаточно вызвать метод model.detect, первым аргументом которого является ссылка на DOM элемент video (строка 14). Этот метод возвращает Promise, результатом которого является массив с метаинформацией о каждом объекте, который был определен моделью в следующем формате:

[{
  bbox: [x, y, width, height],
  class: "person",
  score: 0.8380282521247864
}, {
  bbox: [x, y, width, height],
  class: "kite",
  score: 0.74644153267145157
}]

Функция buildObjectReactngle – отрисовывает области на canvas вокруг объектов, распознанных моделью. Canvas наложен на видео поток сверху.

Здесь ссылка на git-repository с полным кодом.

2. Использование сериализированных обученных моделей TensorFlow

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

В TensorFlow.js в АПИ есть специальный метод tf.GraphModel.save. Модель сохраняется в TensorFlowJS JSON формате. Это набор файлов model.json и один или более бинарных файлов (рисунок 3). Файл model.json содержит информацию о топологии сети и имеет исчерпывающую информации о классах, из которых состоит модель, а также конфигурации слоев модели. Бинарные файлы содержат значения весов всех слоев модели и если модель имеет большое число параметров, то они могут разбиваться на шарды, по умолчания каждый бинарный файл не более 4МБ.

Рисунок 3 – Структура сохраненной модели в формате TensorFlow.js JSON.

На самом деле TensorFlowJS может работать не только с моделями, которые были сохранены им же, но так же есть возможность работать с моделями, которые были обучены с помощью фреймворка Keras на языке Python, но сохраненных с помощью TensorFlow в формате SavedModel. Всего есть 3 типа формата сериализации моделей, которые вы сможете использовать с TensorFlowJS (рисунок 4):

TensorFlow SavedModel – формат модели по умолчанию, с которыми модели сохраняются с помощь TensorFlow;
Keras Model – модели, сохраняемые фреймворком Keras в формате HDF5;
TensorFlow Hub Model – модели, которые распространяются через специальную платформу TensorFlow Hub.

Однако перед использование SavedModel, Keras Model, Tensorflow Hub model в TensorFlowJS необходимо конвертировать эти модели с помощью tfjs_converter.

Тут следует обратить внимание, что модели в форматах SavedModel и TensorFlow Hub могут быть конвертированы только в формат tfjs_graph_model. Модели же, сохраненные в формате Keras, могут конвертированы как в формат tfjs_graph_model, так и в формат tfjs_layers_model.

Чем же отличаются tfjs_layers_model от tfjs_graph_model?

Модели в формате tf_js_graph_model преобразуются в экземпляр класса tf.FrozenModel, что означает что все параметры модели зафиксированы и не подлежат изменению. Таким образом, если вы хотели бы переобучить модель с новой выборкой входных данных, которая более релевантна решаемой вами задачи, или же изменить топологию модели на базе существующей – то вам такой формат модели не подойдет. Однако есть преимущество для этой модели - это что вычисление (inference time) будет значительно меньше, по сравнению с той же моделью, но преобразованной в формат tfjs_layers_model.

Модель в этом формате загружается с помощью следующего АПИ:

const model = tf.loadGraphModel('/path/to/model.json')

Если есть нужда в переобучении модели или изменении ее топологии на стороне браузера, то модель должна быть загружена в формате tfjs_layers_model. Модель в этом формате загружается с помощью следующего АПИ:

const model = tf.loadLayersModel('/path/to/model.json')

На рисунке 5 представлена сводная схема по вышесказанному.

Рисунок 5 – Сводная таблица по конвертации сериализованных моделей

Конвертор можно использовать прямо из командной строки или же его можно вызывать непосредственно в Python-скрипте.

Использование конвертора из командной строки

Для этого вам необходимо установить tensorflowjs:

$ pip  install tensorflowjs

Предположим, мы имеем Keras модель, сохраненная в HDF5 формате в tfjs_layers_model формат, для этого достаточно вызвать команду:

$ tensorflowjs_converter \ 
      --input_format keras \
      --output_format tfjs_layers_model \
      path/to/my_model.h5 \
      path/to/tfjs_target_dir

После работы конвертора, в папке path/to/tfjs_target_dir вы увидите знакомые уже вам файлы model.json и и бинарные файлы содержащие значения весов модели.

Использование конвертора из Python-скрипта

Установите зависимости для скрипта:

$ pipenv keras tensorflowjs

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

import keras
import tensorflowjs as tfjs
mobileNet = keras.applications.mobilenet_v2.MobileNetV2()
tfjs.converters.save_keras_model(mobileNet, './model/from_python_script')

Некоторые ошибки, с которыми вы можете столкнуться

После того, как вы конвертировали модель в tfjs_layers_model и вы пытаетесь загрузить модель на клиенте с помощью tf.loadLayersModel, то можете получить такого рода ошибку:

Uncaught (in promise) Error: Unknown layer: Functional. This may be due to one of the following reasons:
1. The layer is defined in Python, in which case it needs to be ported to TensorFlow.js or your JavaScript code.
2. The custom layer is defined in JavaScript, but is not registered properly with tf.serialization.registerClass().
    at deserializeKerasObject (generic_utils.ts:242)
    at deserialize (serialization.ts:31)
    at loadLayersModelFromIOHandler (models.ts:294)
    at async loadModel (index.js:9)

Если снова взгляните на рисунок 3, то увидите, что файл содержит описание классов, на базе которых строится модель. При этом мы видим что модель создана с помощь класса Functional. Вам надо обратить внимание на версию TF.js, которую вы используете и его АПИ. В данном случае использовался TF.js версии 2.0, и если вы посмотрите на АПИ, то увидите, что в АПИ этой версии нет класса tf.Functional. Именно тут скрывается и проблема. Как его можно разрешить:
- обновить версию TF.js до версии, где в АПИ представлен класс tf.Functional
- в связи с тем, что tf.Functional extends tf.LayersModel, то в model.json файле вместо Functional в поле class_name нужно использовать Model

Также можете посмотреть об ошибке на stackoverflow тут.

Как использовать загруженную модель?

В отличии от упакованных моделей в абстрактный АПИ, тут нам не избежать знакомства с TensorFlow, так как помимо загрузки модели, нам также необходимо позаботится о передачи входных данных в модель.

Итак, давайте просто используем предварительно конвертированную модель MobileNetV2 из формата Keras в формат tfjs_graph_model для классификации изображений.

После загрузки изображения, для того чтобы понять в каком формате нам надо передавать данные в модель, посмотрите файл model.json “inputs” поле:

 "inputs": {
    "input_1:0": {
        "name": "input_1:0",
        "dtype": "DT_FLOAT",
        "tensorShape": {
            "dim": [
                { "size": "-1"},
                { "size": "224"},
                { "size": "224"},
                { "size": "3"}
            ]
        }
    }
},

Таким образом на вход модель необходимо передать 4D-tensor размерностью [null, 224, 224, 3], что соответствует [EXAMPLE_BATCH, WIDTH, HEIGHT, COLOR_CHANNELS]. Вдоль первой оси может располагаться одновременно несколько изображений, но так как мы будем разрабатывать приложение, которое позволяет пользователю загрузить только одно изображение, то в нашем случае на вход модели всегда будем передавать тензор размерностью [1, 224, 224, 3].

Также взглянем, что модель нам выдаст в файле model.json в поле “outputs”:

"outputs": {
    "Identity:0": {
        "name": "Identity:0",
        "dtype": "DT_FLOAT",
        "tensorShape": {
            "dim": [
                { "size": "-1" },
                { "size": "1000" }
            ]
        }
    }
}

Это 2D-тензор размерностью [null, 1000], где 1000 – это количество классов, на которые наша сеть может классифицировать изображения. Это так называемый one-hot вектор, в котором все значения равны нулю, за исключением одного, например [0, 0, 1, 0, 0] – это значит что нейронная сеть считает с вероятностью 1, что на изображении класс с индексом 2 (индексация как обычно начинается с нуля). Когда мы будет использовать модель, то этот вектор будет представлять собой распределение вероятности для каждого из классов, например, можно получить такие значения: [0.07, 0.1, 0.03, 0.75, 0.05]. Обратите внимание, что сумма всех значений будет равна 1, а модель считает с максимальной вероятностью 0.75, что это объект класса c индексом 3.

Теперь все готово. Полный код представлен тут:

import * as tf from '@tensorflow/tfjs';
import {GraphModel, Tensor2D} from '@tensorflow/tfjs';
// other dependencies

export default () => {
    const [model, setModel] = useState<GraphModel>();
    const [image, setImage] = useState<ImageType>();
    const [results, setResults] = useState<ResultType[]>([]);

    useEffect(() => {
        (async () => {
            setModel(await tf.loadGraphModel('/mobileNetV2/model.json'));
        })();
    }, []);

    useEffect(() => {
        tf.tidy(() => {
            if (image?.image && model) {
                (async () => {
                    const offset = tf.scalar(127.5);
                    const input = tf.browser.fromPixels(image.image)
                        // make image compatable with model input
                        .resizeNearestNeighbor([224, 224]) 
                        // feature scale tensor image to range [-1, 1]
                        .sub(offset).div(offset) 
                        .toFloat()
                        // adding batch axis and convert tensor with shape
                        // [224, 224, 3] to [1, 224, 224, 3]
                        .expandDims();
                    const output = model.predict(input) as Tensor2D;
                    const results = Array.from(await output.data())
                        .map((item, i) => ({probability: item, label: labels[i]}))
                        .sort((a1, a2) => a2.probability - a1.probability)
                        .slice(0, 5);
                    setResults(results);
                })();

            }

        });
    }, [image]);

    return (
        <div>
            {model && <InputFile setImage={setImage}/>}
            {image && image.src && <img src={image.src} alt={'Image'}/>}
            {results && results.length > 0 && results.map((item, i) => (
                <div key={i}>
                    {item.label} with probability {item.probability.toFixed(5)}
                </div>
            ))}
        </div>
    );
}

При рендеринге реакт-компонента, мы загружаем модель, которая была конвертирована с Keras (строки 10-14). Как только модель будет загружена – пользователю откроется возможность загружать изображение для классификации.

Как только пользователь выберет изображение – мы для начала должны преобразовать изображения в тензор размерностью [224, 224, 3], а так же необходимо нормализировать данные, чтобы значения пикселей было в промежутке [-1, 1]. Передадим этот тензор модели и получим выходной тензор (строка 30). Теперь нам надо всего лишь сопоставить текстовые метки с индексами, и вывести пять наиболее вероятных классов для загруженного изображения.

Обратите тут внимание, что функция обернута в tf.tidy – это важная деталь. Для увеличения производительности, TensorFlow.js в браузере обычно вычисления производит с помощью WebGL, и к сожалению тут нет механизма сборщика мусора, который автоматически как-то мог бы понять, что некоторые тензоры уже не нужны и можно высвободить память с WebGL. Чтобы этот процесс не был ручным (потому что каждый тензор в своем АПИ имеет метод dispose, который высвобождает память), мы оборачиваем операции над тензором в tf.tidy и после исполнения функции выживут только те тензоры, которые возвращаются этой функцией, а все оставшиеся тензоры будут уничтожены. Так как мы ничего не возвращаем из функции - то все тензоры будут уничтожены после ее исполнения.

Полный код вы можете найти в git-repository тут

TensorFlowJS: использование обученных моделей без их модификаций в браузере