Замена классификационной головы

Keras.js — это JavaScript-библиотека, позволяющая выполнять модели, обученные в Keras (Python), непосредственно в браузере или на Node.js без необходимости серверной обработки. Основное преимущество заключается в переносимости нейронных сетей и возможности работы с ними на клиентской стороне, что снижает задержки при прогнозировании и исключает передачу данных на сервер.

Модели Keras.js используют формат JSON для описания архитектуры и веса модели, экспортированные из Keras через model.to_json() и model.save_weights(). После загрузки они становятся доступными для выполнения предсказаний через API Keras.js.

Структура модели и ключевые компоненты

Модель в Keras.js состоит из двух частей:

  1. Архитектура модели (JSON) — содержит информацию о слоях, их типах, функциях активации, размере входа и выхода.
  2. Веса модели (binary .bin файлы) — числовые значения параметров сети, загруженные отдельно для ускорения передачи и загрузки.

Для выполнения предсказаний используется объект KerasJS.Model, который инициализируется указанием пути к JSON и весам.

const model = new KerasJS.Model({
  filepaths: {
    model: 'model.json',
    weights: 'model_weights.buf'
  },
  gpu: true
});
await model.ready();

Замена классификационной головы

Классификационная голова (classification head) — это верхний слой нейронной сети, который отвечает за формирование окончательного прогноза. В типичных CNN она состоит из одного или нескольких плотных (Dense) слоев, завершающихся softmax для многоклассовой классификации или sigmoid для бинарной.

Замена головы необходима при адаптации предобученной модели к новой задаче, например, если сеть обучена на ImageNet (1000 классов), а требуется классификация на 10 классов.

Шаг 1. Загрузка предобученной модели

Предварительно экспортированная модель загружается в Keras.js так же, как любая другая:

const baseModel = new KerasJS.Model({
  filepaths: {
    model: 'base_model.json',
    weights: 'base_model_weights.buf'
  },
  gpu: true
});
await baseModel.ready();

Шаг 2. Извлечение и модификация архитектуры

Архитектура модели хранится в JSON. Для замены головы необходимо:

  1. Скопировать все слои до классификационной головы.
  2. Удалить исходный верхний слой или блок слоев.
  3. Добавить новый слой Dense с количеством нейронов, соответствующим новой задаче.

Пример структуры нового слоя:

{
  "class_name": "Dense",
  "config": {
    "units": 10,
    "activation": "softmax"
  },
  "name": "new_classification_head"
}

После модификации JSON его можно сохранить и загрузить в Keras.js так же, как исходную модель.

Шаг 3. Подготовка весов для нового слоя

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

Шаг 4. Формирование новой модели

После изменения архитектуры создается новый объект KerasJS.Model, указывающий на измененный JSON и веса:

const newModel = new KerasJS.Model({
  filepaths: {
    model: 'modified_model.json',
    weights: 'modified_model_weights.buf'
  },
  gpu: true
});
await newModel.ready();

Шаг 5. Инференс с новой головой

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

const inputData = new Float32Array([/* данные */]);
const outputData = await newModel.predict({ input: inputData });
console.log(outputData.output);

Важные моменты при замене головы

  • Совместимость размеров входа и выхода: Новый слой должен иметь количество нейронов, соответствующее числу классов новой задачи.
  • Активация: Для многоклассовой классификации используется softmax, для бинарной — sigmoid.
  • Инициализация весов: Старые веса не подходят для нового слоя, их нужно генерировать заново.
  • GPU vs CPU: При больших моделях рекомендуется включать GPU-режим для ускорения предсказаний.

Расширенные сценарии

  • Множественные головы: Для задач с несколькими выходами можно добавить несколько отдельных Dense слоев, подключенных к общей базе.
  • Фризинг базовых слоев: Для тонкой настройки достаточно заморозить веса всех слоев до классификационной головы и обучать только новый слой на небольшом наборе данных.
  • Сохранение новых весов: После обучения на новом датасете веса можно экспортировать в .buf для повторного использования в Keras.js.

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