Заморозка базовых слоёв

Заморозка слоёв (Layer Freezing) — ключевой приём при работе с предобученными моделями в TensorFlow.js, позволяющий использовать уже выученные признаки и предотвращать их изменение во время дообучения. Эта техника особенно актуальна при переносе знаний (transfer learning), когда модель обучается на новой задаче с ограниченным объёмом данных.

Принцип работы

Предобученные модели содержат слои, параметры которых уже оптимизированы под определённые задачи, например классификацию изображений. Заморозка слоёв предотвращает обновление весов во время оптимизации, оставляя их статичными. Это достигается установкой свойства trainable в значение false:

layer.trainable = false;

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

Практическая реализация

В TensorFlow.js предобученные модели можно загрузить через API tf.loadLayersModel или использовать встроенные модели из @tensorflow-models. Например, для работы с MobileNet:

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

async function createModel() {
  const mobilenetModel = await mobilenet.load();
  
  // Получаем базовую модель без верхних слоёв
  const baseModel = mobilenetModel.model;
  for (const layer of baseModel.layers) {
    layer.trainable = false; // Заморозка всех базовых слоёв
  }

  // Добавляем новые слои для специфической задачи
  const model = tf.sequential();
  model.add(baseModel);
  model.add(tf.layers.flatten());
  model.add(tf.layers.dense({ units: 128, activation: 'relu' }));
  model.add(tf.layers.dense({ units: 10, activation: 'softmax' }));

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

  return model;
}

В этом примере заморожены все слои MobileNet, а новые полносвязные слои адаптированы под задачу классификации на 10 классов.

Выбор слоёв для заморозки

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

// Заморозка первых 100 слоёв, остальные обучаемы
baseModel.layers.slice(0, 100).forEach(layer => layer.trainable = false);

Особенности компиляции

После изменения свойства trainable необходимо перекомпилировать модель:

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

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

Эффекты и преимущества

  1. Снижение вычислительной нагрузки: замороженные слои не требуют расчёта градиентов.
  2. Быстрое обучение: обучение новых слоёв происходит быстрее, поскольку базовые признаки уже извлечены.
  3. Улучшенная стабильность: предотвращается разрушение уже выученных признаков.
  4. Меньший риск переобучения: количество обучаемых параметров сокращается.

Советы по использованию

  • При небольшом объёме данных лучше замораживать большую часть модели.
  • Для крупных наборов данных допускается разморозка некоторых высокоуровневых слоёв для дообучения.
  • Комбинация заморозки слоёв и регуляризации (Dropout, L2) помогает достичь лучшей генерализации.

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

Оглавление