Форматы сохранения: JSON + бинарные веса

TensorFlow.js предоставляет мощные возможности для работы с моделями машинного обучения прямо в браузере или на сервере с Node.js. Одной из ключевых задач является сохранение и восстановление обученных моделей для последующего использования без необходимости повторного обучения. Модель в TensorFlow.js состоит из двух основных компонентов:

  1. Архитектура модели — описывает структуру сети, слои, параметры активации, связи между слоями.
  2. Веса модели — значения весов и смещений нейронных сетей, полученные после обучения.

Для оптимального управления этими компонентами используется формат JSON + бинарные веса.


Формат JSON для описания архитектуры

Архитектура модели сохраняется в виде JSON-объекта. Основные элементы структуры:

  • modelTopology — описание слоев и их параметров. Включает информацию о типах слоев, их конфигурации и порядке соединений.
  • trainingConfig — параметры обучения модели, если сохранение происходит после обучения (optimizer, loss, metrics).
  • weightsManifest — описание файлов с бинарными весами, их размеры и типы данных.

Пример структуры JSON:

{
  "modelTopology": {
    "class_name": "Sequential",
    "config": {
      "layers": [
        {
          "class_name": "Dense",
          "config": {
            "units": 128,
            "activation": "relu",
            "inputShape": [784]
          }
        },
        {
          "class_name": "Dense",
          "config": {
            "units": 10,
            "activation": "softmax"
          }
        }
      ]
    }
  },
  "weightsManifest": [
    {
      "paths": ["weights.bin"],
      "weights": [
        {
          "name": "dense/kernel",
          "shape": [784, 128],
          "dtype": "float32"
        },
        {
          "name": "dense/bias",
          "shape": [128],
          "dtype": "float32"
        }
      ]
    }
  ]
}

Ключевые моменты:

  • JSON-файл сам по себе не содержит численные значения весов.
  • Он хранит только метаданные и описание модели.
  • Формат позволяет легко модифицировать архитектуру без влияния на бинарные данные весов.

Бинарные веса модели

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

  • Данные сохраняются в виде последовательности Float32 или Int32.
  • Каждое значение соответствует параметру конкретного слоя (ядро или смещение).
  • Файл может содержать несколько блоков, если веса разбиты по слоям или файлам.

Преимущества бинарного формата:

  • Скорость загрузки — браузер или Node.js может напрямую интерпретировать бинарные данные.
  • Экономия памяти — по сравнению с текстовыми форматами JSON или CSV.
  • Совместимость — веса легко интегрируются с TensorFlow.js и TensorFlow (Python) через конвертеры.

Методы сохранения моделей

TensorFlow.js предоставляет встроенные методы для сохранения моделей с разделением на JSON и бинарные веса:

// Сохранение модели в файловую систему (Node.js)
await model.save('file://./my-model');

// Сохранение модели в браузере через IndexedDB
await model.save('indexeddb://my-model');

// Сохранение модели через HTTP POST (сервера или облако)
await model.save('http://example.com/upload-model');

Особенности:

  • model.save() автоматически формирует JSON-файл для архитектуры и отдельный бинарный файл для весов.
  • В браузере можно использовать localstorage или indexeddb для хранения модели между сессиями.
  • В Node.js возможна прямая запись файловой структуры на диск.

Загрузка модели

Восстановление модели из формата JSON + бинарные веса выполняется с помощью метода tf.loadLayersModel:

// Загрузка модели из файловой системы (Node.js)
const model = await tf.loadLayersModel('file://./my-model/model.json');

// Загрузка модели из IndexedDB (браузер)
const model = await tf.loadLayersModel('indexeddb://my-model');

// Загрузка модели с удалённого сервера
const model = await tf.loadLayersModel('https://example.com/my-model/model.json');

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

  • JSON-файл парсится для восстановления структуры модели.
  • Бинарные файлы весов загружаются отдельно.
  • Все слои и веса восстанавливаются в памяти, готовые к использованию.

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

  1. Разделение веса и архитектуры — позволяет обновлять веса без изменения структуры и наоборот.
  2. Использование нескольких файлов весов — улучшает производительность при больших моделях.
  3. Контроль типов данных — весовые коэффициенты обычно хранятся как float32; несоответствие типов может вызвать ошибки при загрузке.
  4. Версионирование моделей — рекомендуется включать версию в JSON для совместимости при обновлениях архитектуры.

Особенности работы с большими моделями

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

Интеграция с конвертерами TensorFlow

Модели, обученные в Python TensorFlow, могут быть конвертированы в TensorFlow.js формат (JSON + бинарные веса) с помощью утилиты tensorflowjs_converter. Это обеспечивает:

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

Система сохранения JSON + бинарные веса в TensorFlow.js предоставляет надежный, гибкий и высокопроизводительный способ управления моделями, позволяя эффективно работать как с небольшими экспериментальными сетями, так и с крупными продвинутыми нейросетями.