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

Кроме того, на мой взгляд Transfer Learning – это наиболее перспективная техника для использования на клиентской стороне с помощью TensorFlowJS. Большим преимуществом тут перед применением той же самой техники на сервере – это сохранение конфиденциальности клиентской информации и наличием возможности доступа к сенсорам (камера, гео-локации и др).

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

Рисунок 1 – Структурная схема Transfer Learning

Для загрузки модели в TensorFlow используют специальный JSON формат, который в свою очередь может быть двух типов: graph-model или layers-model (рисунок 2).

Рисунок 2 – Классификация форматов сериализации моделей, совместимые с TensorFlowJS

Для Transfer Learning модель в формате graph-model не применима, мы можем использовать эту модель только в том виде, в котором ее загрузили, без возможности ее переобучении на новой выборке данных или же изменении топологии под свои нужды. Поэтому ниже мы будет вести речь, только про модель в формате layers-model.

Для загрузки предварительно обученной модели в формате layers-model, TensrFlowJS предоставляет в своем АПИ метод tf.loadLayersModel , который загружает топологию модели и ее веса, полученные в результате длительного процесса обучения. Топология модели задаётся в json формате, который также содержит поле weightsManifest с указанием путей в бинарном формате, содержащих все веса связей обученной нейронной сети (в случае сложной нейронной сети с миллионами обучаемыми параметров, веса модели могут представлены в нескольких файлах-шардах).

С помощью специального конвертора вы можете конвертировать модели, например обученных на Python с помощью фреймворка Keras (он сейчас включен как подмодуль в TensorFlow) в формат, совместимый с TF.js. Более подробно вы можете узнать тут.

После загрузки модели, мы можем модифицировать и переобучить ее согласно требованиям задачи. Так как мы будем стремиться уменьшить время на обучение модели на базе новой выборки данных, то очевидным будет то, что необходимо “замораживать” (freeze) как можно больше слоев в исходной модели. Заморозить слой – это перевод всех нейронов слоя из категории обучаемых (trainable) в категорию необучаемых (untrainable). Во время обучения сети, оптимизатор будет модифицировать только обучаемые параметры сети, что может значительно сократить время на ее обучение.

Справедливым вопросом тогда будет – как много и какие слои мы должны замораживать. Чтобы интуитивно понимать это, давайте рассмотрим типовую структурную схему нейронной сети, которая связана с классификаций изображений между N классами. Структурная схема приведена на рисунке 3.

Рисунок 3 – Топология нейронной сети для классификации изображений между N классами

Нейронная сеть включает в себя несколько последовательно подключенных сверточных слоев (convolutional layers) и в конце подключено несколько скрытых полносвязных слоев нейронных сетей (dense layers, fully-connected layers).

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

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

Рисунок 4 - Признаки, выделяемые сверточными слоями на разных уровнях нейронной сети

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

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

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

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

!!!! Обобщим, выше сказанное – чем больше у нас обучающая выборка для обучающей модели, тем меньше слоев нам надо замораживать.

Визуализация алгоритма, описанного выше, вы можете посмотреть на рисунке 5.

Рисунок 5 – Схема настройки новой нейронной сети на базе уже обученной

Практика

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

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

Рисунок 6 – Структурная диаграмма по созданию новой модели на базе обученной модели MobileNet

1 шаг – Загрузка обученной модели MobileNet и ее анализ
MobileNet – модель, обученная на базе ImageNet выборке (коллекция, включающая более миллиона изображений, разбитую вручную между 1000 классами).

Загрузим модель с помощью специального пользовательского React-hook:

export default () => {
    const [model, setModel] = useState<LayersModel>();

    useEffect(() => {
        (async function init() {
            const model = await tf.loadLayersModel(MODEL_URL);
            setModel(() => model);
            model.summary();
        })();
    }, []);
    
    return model;
}

Посмотрим на топологию загруженной сети вызовом model.summary(), рисунок 7.

Рисунок 7 – Топология загруженной модели

Модель состоит из 88 слоев, входной слой которого представлен тензором размерностью [null, 224, 224, 3], принимающий изображение размерностью 224x224 пикселей с тремя цветовыми каналами. Выходной слой – это тензор размерность [null,1000].

Выходной тензор [null, 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.

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

Слои модели можно разбить условно на слои, отвечающих за:
- извлечение характерных признаков изображения
- классификацию изображений между 1000 классами
Для новой модели, мы хотим использовать все слои, отвечающих за извлечение признаков изображений, но при этом исключим все слои, отвечающих за классификацию изображений между 1000 классами, и вместо этого предоставить свой классификатор, который будет производить классификацию между 3мя классами.

На рисунке 8 показана диаграмма разбивка загруженных слоев по их назначению (рисунок 8).

Рисунок 7 – Создание новой модели на базе существующей

Кстати, тут можно проанализировать нашу задачу на предмет, стоит ли замораживать слои. Во-первых, однозначно мы не будем заставлять пользователя делать большую тренировочную выборку для обучения сети, будет вполне себе достаточно по 30-50 изображений каждого класса. Также решаемая нами задача – классификация между 3 классами изображений с жестами, достаточно близка к задаче классификации изображений между 1000 классами, которая была предварительно обучена на большом наборе тренировочных изображений. Следовательно, согласно рисунку 5, мы можем заморозить все слои части модели MobileNetV2, которые мы извлекли для своих нужд:

// freeze all layers of MobileNet
for (const layer of pretrainedModel.layers) {
    layer.trainable = false;
}

Шаг 2 – Поиск последнего слоя в загруженной модели, отвечающего за извлечение характерных признаков изображения

Для этого будем использовать АПИ tf.LayersModel.getLayer, который позволяет найти искомый слой либо по его индексу, либо по уникальному имени. Более предпочтительным является поиск слоя по уникальному имени, чем по его индексу. Обратите внимание (рисунок 7), каждый слой загруженной модели имеет уникальное название: input_1; conv_pw_13_bn; conv_pw_13_relu и др. Как и договаривались – найдем последний сверточный слой в модели и это является слой с именем conv_pw_13_relu. Таким образом код будет выглядеть следующим образом:

const truncatedLayer = pretrainedModel.getLayer('conv_pw_13_relu');
const truncatedLayerOutput = truncatedLayer.output as SymbolicTensor;

Шаг 3 – Создание нового классификатора для классификации между 3 классами

Модель будет состоять из нескольких слоев (рисунок 8):
- первый слой - мы должны многомерный тензор, который получим от сверточных слоев MobileNetV2 преобразовать в одномерный массив, чтобы он был совместим с полносвязными (fully-connected, dense layers) слоями модели для классификации;

- второй слой – скрытый полносвязный слой, который будет содержать 100 нейронов (параметр может быть подобран на базе экспериментов, это так называемый hyper parameter) с активационной функцией RELU. RELU – наиболее распространенная активационная функция, которая добавляет нелинейность модели;

- третий слой – выходной полносвязный слой, содержащий ровно 3 нейрона, каждый из которых определяет вероятность принадлежности изображения к одному из трех классов: камень, ножницы, бумага. Слой будет с активационной функцией SOFTMAX. SOFTMAX функция, которая всегда используется для последнего слоя для решения задач множественной классификации. Она преобразует вектор z размерности K в вектор той же размерностью σ, где каждое значение представлена вещественным числом в интервале [0, 1], а их сумма равна единице. И как уже указывалось выше,каждое значение в векторе σ_i трактуются как вероятность того, то объект принадлежит классу i.

Рисунок 8 – Новые слои для модели, отвечающие за классификацию между 3 классами

Код будет выглядеть следующим образом:

function buildNewHead(inputShape: Shape, numClasses: number) {
    // Creates a 2-layer fully connected model
    return tf.sequential({
        layers: [
            tf.layers.flatten({
                name: 'flatten',
                inputShape: inputShape.slice(1),
            }),
            tf.layers.dense({
                name: 'hidden_dense_1',
                units: 100,
                activation: 'relu'
            }),
            // Layer 2. The number of units of the last layer should 
            // correspond to the number of classes we want to predict.
            tf.layers.dense({
                name: 'softmax_classification',
                units: numClasses,
                activation: 'softmax'
            })
        ]
    });
}

Шаг 4 - Связывание нового классификатора со слоями предварительно обученной модели

В шаге 2 мы нашли последний слой в графе модели, который мы хотим оставим в новой разрабатываемой модели. Давайте свяжем этот слой с классификатором, созданным на шаге 3 (см. рисунок 6). Для этого экземпляры класса модели/слоев имеют специальный метод apply, который делает соединение между слоями. Это будет выглядеть следующим образом:

const transferHead = buildNewHead(truncatedLayerOutput.shape, numClasses);
const newOutput = transferHead.apply(truncatedLayerOutput) as SymbolicTensor;

Шаг 5 – Создание новой модели и ее компиляция

Для создания новой модели на базе существующей мы будет использовать tf.model, который принимает в качестве аргумента объект с двумя обязательными полями inputs и outputs. Оба параметра ожидают массив, где каждый элемент массива представляет собой Symbolic tensor (отличие от обычного тензора – это то, что по сути это описание тензора имеющая сведения о размерности и типе, но при этом отсутствуют данные, так называемый placeholder).

Каждый слой в модели имеет два атрибута – вход и выход, представленные Symbolic tensor.

Тут иногда удобно сопоставлять слой с обычной JS функцией, которая принимает входной параметр, внутри производится некоторая логика по вычислению и всегда возвращает выходной параметр. Отличие тут, что входными и выходными параметрами будут Symbolic Tensor.

Итак, нам надо определить inputs и outputs для новой модели. Input для новой модели будет совпадать с input загруженной модели и доступно через: pretrainedModel.inputs.

А выходным Symbolic Tensor – это выход после последнего слоя нового классификатора, который доступен в переменной newOutput, после применения метода apply на предыдущем шаге.

// build new model
const model = tf.model({
    inputs: pretrainedModel.inputs,
    outputs: newOutput
});

Скомпилируем модель, задав нужный оптимизатор и loss-функцию. Наиболее лучшим оптимизатором является tf.train.adam, а loss-функция для классификации задают обычно categoricalCrossentropy:

model.compile({
    optimizer: tfC.train.adam(0.0001),
    loss: 'categoricalCrossentropy'
});

Полный код создания новой модели на базе предварительно обученной MobileNet

import {useEffect, useState} from 'react';
import {LayersModel, Shape, SymbolicTensor} from '@tensorflow/tfjs';
import * as tf from '@tensorflow/tfjs';

function buildNewHead(
      inputShape: Shape, 
      numClasses: number): LayersModel {
    // Creates a 2-layer fully connected model
    return tf.sequential({
        layers: [
            tf.layers.flatten({
                name: 'flatten',
                inputShape: inputShape.slice(1),
            }),
            tf.layers.dense({
                name: 'hidden_dense_1',
                units: 100,
                activation: 'relu',
                kernelInitializer: 'varianceScaling',
                useBias: true
            }),
            // Layer 2. The number of units of the last layer 
            // should correspond to the number of classes 
            // we want to predict.
            tf.layers.dense({
                name: 'softmax_classification',
                units: numClasses,
                kernelInitializer: 'varianceScaling',
                useBias: false,
                activation: 'softmax'
            })
        ]
    });
}

export default (
    pretrainedModel: LayersModel | undefined, 
    numClasses: number) => {
    const [model, setModel] = useState<LayersModel>();

    useEffect(() => {
        if (pretrainedModel) {
            // find last convolutional layer by its name
            // it is a layer followed by layers which are responsible
            // for classification
            const truncatedLayer = 
                  pretrainedModel.getLayer('conv_pw_13_relu');
            const truncatedLayerOutput = 
                  truncatedLayer.output as SymbolicTensor;

            // freeze all layers of MobileNet
            for (const layer of pretrainedModel.layers) {
                layer.trainable = false;
            }

            const transferHead = buildNewHead(
                  truncatedLayerOutput.shape, numClasses);
            const newOutput = 
                  transferHead.apply(truncatedLayerOutput) as SymbolicTensor;

            const model = tf.model({
                inputs: pretrainedModel.inputs,
                outputs: newOutput
            });

            model.compile({
                optimizer: tf.train.adam(0.0001),
                loss: 'categoricalCrossentropy'
            });

            setModel(model);
        }
    }, [pretrainedModel]);

    return model;
}

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

function App() {
    // skipped for brevity
    const train = async () => {
        if (model
            && controllerDataset.current.xs !== null
            && controllerDataset.current.ys !== null) {
            const {current: {xs, ys}} = controllerDataset;
            model.fit(xs, ys, {
                batchSize: 24,
                epochs: 20,
                shuffle: true,
                callbacks: {
                    onBatchEnd: async (batch, logs) => {
                        setLoss(() => (logs?.loss as number).toFixed(5));
                    },
                    onTrainEnd: () => {
                        changeTrainingState(() => TRAINING_STATES.TRAINED);
                    }
                }
            });
        }
    };
   
    return (
        <div>
            {/* skipped for brevity*/}
             <button 
                  disabled={trainingState === TRAINING_STATES.TRAINING}
                  onClick={() => {
                      if (trainingState !== TRAINING_STATES.TRAINING) {
                         changeTrainingState(() => TRAINING_STATES.TRAINING);
                         train();
                       }
              }}>
                 Train
              </button>
        </div>
    );
}

После того, как процесс обучения будет закончить, необходимо запустит процесс использования этой модели для классификации изображений:

const predict = async () => {
    if (webcamIterator && model) {
        const image = await webcamIterator.capture();
        const processedImage = tf.tidy<Tensor4D>(
            () => image.expandDims().toFloat().div(127).sub(1)
        );
        const label = 
               tf.argMax(model.predict(processedImage) as Tensor, 1)
                    .dataSync();
        setActiveLabel(() => label[0]);
        image.dispose();
        processedImage.dispose();
        requestAnimationFrame(() => predict());
    }
};

useEffect(() => {
    if (trainingState === TRAINING_STATES.TRAINED) {
        requestAnimationFrame(() => predict());
    }
}, [trainingState]);

Полный код React-component

import React, {useEffect, useRef, useState} from 'react';
import useUploadModel from './hook/useUploadModel';
import useWebcamIterator from './hook/useWebcamIterator';
import WorkingArea from './components/WorkingArea';
import useNewModelTopology from './hook/useNewModelTopology';
import ControllerDataset from './service/ControllerDataset';
import * as tf from '@tensorflow/tfjs';
import {Tensor4D, Tensor} from '@tensorflow/tfjs';
import {Container, Button, Grid, Box, TextField} from '@material-ui/core';

const NUM_CLASSES = 3;

enum TRAINING_STATES {
    NOT_TRAINED,
    TRAINING,
    TRAINED
}

function App() {
    const uploadedModel = useUploadModel();
    const model = useNewModelTopology(uploadedModel, NUM_CLASSES);
    const {videoRef, webcamIterator} = useWebcamIterator();

    const [loss, setLoss] = useState<string | null>(null);
    const controllerDataset = 
               useRef(new ControllerDataset(NUM_CLASSES));

    const [trainingState, changeTrainingState] = 
               useState<TRAINING_STATES>(TRAINING_STATES.NOT_TRAINED);
    const [activeLabel, setActiveLabel] = useState<number>();

    const train = async () => {
        if (model
            && controllerDataset.current.xs !== null
            && controllerDataset.current.ys !== null) {
            const {current: {xs, ys}} = controllerDataset;
            model.fit(xs, ys, {
                batchSize: 24,
                epochs: 20,
                shuffle: true,
                callbacks: {
                    onBatchEnd: async (batch, logs) => {
                        setLoss(() => (logs?.loss as number).toFixed(5));
                    },
                    onTrainEnd: () => {
                        changeTrainingState(
                              () => TRAINING_STATES.TRAINED
                        );
                    }
                }
            });
        }
    };

    const predict = async () => {
        if (webcamIterator && model) {
            const image = await webcamIterator.capture();
            const processedImage = tf.tidy<Tensor4D>(
                () => image.expandDims().toFloat().div(127).sub(1)
            );
            const label = 
                  tf.argMax(model.predict(processedImage) as Tensor, 1)
                       .dataSync();
            setActiveLabel(() => label[0]);
            image.dispose();
            processedImage.dispose();
            requestAnimationFrame(() => predict());
        }
    };

    useEffect(() => {
        if (trainingState === TRAINING_STATES.TRAINED) {
            requestAnimationFrame(() => predict());
        }
    }, [trainingState]);

    return (
        <>
            <video ref={videoRef} width={224} height={224}/>
            <Button
                color="primary"
                variant="contained"
                disabled={trainingState === TRAINING_STATES.TRAINING}
                onClick={() => {
                    if (trainingState !== TRAINING_STATES.TRAINING) {
                        changeTrainingState(
                             () => TRAINING_STATES.TRAINING
                        );
                        train();
                    }
                }}>
                Train
            </Button>
            {loss && <div>Loss: {loss}</div>}
            <WorkingArea numClasses={NUM_CLASSES}
                         webcam={webcamIterator}
                         activeLabel={activeLabel}
                         controllerDataset={controllerDataset.current}
            />
        </>
    );
}

Весь код вы можете изучить в репозитории

Также вы можете изучить следующие статьи:

Transfer Learning с использованием TensorFlow.JS