Построение классификатора поверх готового бэкбона

TensorFlow.js предоставляет возможность создавать сложные модели машинного обучения прямо в браузере или на Node.js, используя готовые предварительно обученные модели как бэкбон (backbone). Бэкбон выполняет роль извлечения признаков из данных, а поверх него строится классификатор, который адаптируется под конкретную задачу.

Выбор и загрузка бэкбона

Наиболее часто используемые бэкбоны в TensorFlow.js — это MobileNet, ResNet, EfficientNet. Они обладают готовыми весами, обученными на ImageNet, что позволяет сразу использовать их для извлечения признаков изображений.

Пример загрузки MobileNet:

import * as tf from '@tensorflow/tfjs';
import * as mobilenet from '@tensorflow-models/mobilenet';

const loadBackbone = async () => {
  const mobilenetModel = await mobilenet.load({ version: 2, alpha: 1.0 });
  const layer = mobilenetModel.model.getLayer('global_average_pooling2d_1');
  const backbone = tf.model({ inputs: mobilenetModel.model.inputs, outputs: layer.output });
  return backbone;
};

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

Подготовка данных

Данные для классификатора должны быть приведены к формату, подходящему для выбранного бэкбона. Для изображений это обычно нормализация пикселей в диапазон [0, 1] или [−1, 1], а также приведение размеров к фиксированным значениям, соответствующим входу модели.

const preprocessImage = (imageElement) => {
  return tf.tidy(() => {
    let tensor = tf.browser.fromPixels(imageElement).toFloat();
    tensor = tf.image.resizeBilinear(tensor, [224, 224]);
    tensor = tensor.div(255.0);
    return tensor.expandDims();
  });
};

Важное правило: использование tf.tidy() помогает управлять памятью и избегать утечек при работе с большим числом изображений.

Создание классификатора поверх бэкбона

После извлечения признаков с помощью бэкбона строится собственный классификатор. Обычно это полносвязные слои с функцией активации relu и на выходе — слой с softmax для многоклассовой классификации.

const buildClassifier = (backbone, numClasses) => {
  const model = tf.sequential();
  model.add(backbone);
  model.add(tf.layers.dense({ units: 128, activation: 'relu' }));
  model.add(tf.layers.dropout({ rate: 0.5 }));
  model.add(tf.layers.dense({ units: numClasses, activation: 'softmax' }));
  return model;
};

Совет: слой dropout помогает уменьшить переобучение при малом объёме данных. Размер скрытого слоя можно подбирать экспериментально, обычно 128–512 нейронов для небольших наборов данных.

Компиляция и выбор оптимизатора

Модель компилируется с функцией потерь и оптимизатором. Для многоклассовой классификации используется categoricalCrossentropy, а для бинарной — binaryCrossentropy.

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

Подсказка: низкая скорость обучения (learningRate) рекомендуется при дообучении бэкбона, чтобы не разрушить предварительно обученные веса.

Дообучение (Fine-tuning)

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

// Заморозка всех слоев бэкбона
backbone.layers.forEach(layer => layer.trainable = false);

// После обучения классификатора
backbone.layers.slice(-10).forEach(layer => layer.trainable = true);

Рекомендация: fine-tuning верхних слоев эффективен, если новая задача сильно отличается от оригинальной (например, классификация медицинских снимков вместо обычных объектов ImageNet).

Обучение модели

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

await model.fit(trainData, {
  epochs: 20,
  validationData: valData,
  batchSize: 32,
  callbacks: tf.callbacks.earlyStopping({ monitor: 'val_loss', patience: 3 })
});

Ключевой момент: использование EarlyStopping предотвращает переобучение, останавливая обучение при отсутствии улучшений на валидации.

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

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

const predict = async (imageElement) => {
  const input = preprocessImage(imageElement);
  const prediction = model.predict(input);
  const predictedClass = prediction.argMax(-1).dataSync()[0];
  return predictedClass;
};

Сохранение и загрузка модели

TensorFlow.js позволяет сохранять модели локально или на сервере для последующего использования без повторного обучения.

await model.save('localstorage://my-model');
const loadedModel = await tf.loadLayersModel('localstorage://my-model');

Совет: при работе в браузере IndexedDB и LocalStorage удобны для небольших моделей, для больших моделей лучше использовать серверное хранилище.

Практические рекомендации

  • Всегда использовать нормализацию входных данных в соответствии с требованиями бэкбона.
  • При малых наборах данных эффективнее сначала тренировать только классификатор, а затем размораживать верхние слои бэкбона.
  • Контролировать потребление памяти в браузере, активно используя tf.tidy() и dispose() для больших тензоров.
  • Экспериментировать с архитектурой классификатора, количеством слоев и нейронов, чтобы достичь оптимального соотношения точности и скорости обучения.

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

Оглавление