Прунинг и сжатие

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

Прунинг весов

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

Типы прунинга:
1. Статический (one-shot) — веса обнуляются единожды после обучения.
2. Динамический (iterative) — веса постепенно обнуляются на протяжении дообучения, что позволяет модели адаптироваться к уменьшенной структуре.
Критерии выбора весов для обнуления:
- Абсолютное значение веса: веса с минимальной величиной считаются наименее значимыми.
- Градиентная чувствительность: веса, изменение которых мало влияет на функцию потерь, являются кандидатами для обнуления.

В TensorFlow.js можно реализовать прунинг с помощью перебора весов модели через API tf.layers или tf.Model. Для каждого слоя извлекаются веса методом getWeights(), затем выполняется маскирование весов с минимальными значениями:

const pruneWeights = (weights, threshold) => {
  return weights.map(w => {
    const data = w.arraySync();
    const prunedData = data.map(v => Math.abs(v) < threshold ? 0 : v);
    return tf.tensor(prunedData, w.shape);
  });
};

const layer = model.layers[0];
const prunedWeights = pruneWeights(layer.getWeights(), 0.01);
layer.setWeights(prunedWeights);

Квантование

Квантование — процесс преобразования весов и активаций модели из формата с плавающей запятой (float32) в более компактные форматы (int8, float16). Квантование уменьшает размер модели и ускоряет вычисления на устройствах с ограниченными ресурсами.

Методы квантования:
- Пост-тренировочное (post-training quantization) — применяется после обучения модели.
- Квантование с обучением (quantization-aware training) — учитывает пониженную точность во время обучения, повышая точность после преобразования.

В TensorFlow.js можно использовать утилиту tfjs-converter для пост-тренировочного квантования при конвертации моделей из формата TensorFlow SavedModel:

tensorflowjs_converter \
    --input_format=tf_saved_model \
    --quantization_bytes=1 \
    /saved_model_path \
    /tfjs_model_path

Сжатие структуры модели

Сжатие может касаться не только весов, но и самой архитектуры сети:

Удаление лишних слоев: слои, которые не сильно влияют на точность, можно убрать.
Снижение количества нейронов в слоях: уменьшение числа нейронов в скрытых слоях сокращает размер модели.
Сжатие с помощью Knowledge Distillation: использование более крупной модели для обучения меньшей модели, которая сохраняет точность, но имеет меньший размер.

Маскирование и сохранение модели

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

const applyMask = (weights, mask) => {
  return weights.mul(mask);
};

После прунинга и сжатия модель можно сохранить для веб-приложений:

await model.save('downloads://pruned_model');

Это создает .json и бинарные .bin файлы, оптимизированные для загрузки и использования в браузере.

Влияние на точность и производительность

Снижение размера модели напрямую уменьшает время загрузки и требования к памяти.
Прунинг при умеренном уровне (10–30% весов) обычно не вызывает сильной деградации точности.
Квантование может привести к небольшой потере точности, особенно при агрессивном переходе к int8.
Комбинация методов (прунинг + квантование) позволяет создавать компактные, быстрые модели для браузеров и мобильных устройств.