Конвертация Python-моделей через tensorflowjs_converter

TensorFlow.js предоставляет возможность использовать обученные модели Python напрямую в среде JavaScript, включая браузер и Node.js. Для этого используется утилита командной строки tensorflowjs_converter, которая выполняет преобразование моделей в формат, совместимый с TensorFlow.js.


Установка tensorflowjs_converter

Утилита поставляется вместе с пакетом tensorflowjs. Для установки используется стандартный менеджер пакетов Python pip:

pip install tensorflowjs

После установки доступна команда tensorflowjs_converter, которая поддерживает конвертацию моделей, сохранённых в формате TensorFlow SavedModel, Keras HDF5 или замороженных графов (Frozen Graph).


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

  1. Keras HDF5 (.h5)

    • Модели, сохранённые с помощью model.save('model.h5').
    • Удобны для быстрого конвертирования в JavaScript.
  2. SavedModel

    • Рекомендуемый формат TensorFlow 2.x.
    • Содержит структуру модели, веса и информацию о сигнатурах.
    • Поддерживается конвертацией с сохранением метаданных.
  3. Frozen Graph (.pb)

    • Статическая графовая модель TensorFlow 1.x.
    • Требует указания входных и выходных узлов при конвертации.

Основные команды tensorflowjs_converter

Конвертация Keras HDF5 в TensorFlow.js

tensorflowjs_converter \
    --input_format=keras \
    path/to/model.h5 \
    path/to/tfjs_model
  • --input_format=keras — указывает формат исходной модели.
  • path/to/model.h5 — путь к исходной Keras-модели.
  • path/to/tfjs_model — директория для сохранения модели в формате TensorFlow.js.

После конвертации в указанной папке создаются файлы:

  • model.json — описывает архитектуру модели и метаданные.
  • group1-shard\*of\*bin — бинарные файлы с весами.

Конвертация SavedModel в TensorFlow.js

tensorflowjs_converter \
    --input_format=tf_saved_model \
    --saved_model_tags=serve \
    path/to/saved_model \
    path/to/tfjs_model
  • --saved_model_tags=serve — тег, соответствующий графу для инференса.
  • Поддерживает указание конкретного подпакета с тегами, если SavedModel содержит несколько вариантов.

Конвертация замороженного графа TensorFlow 1.x

tensorflowjs_converter \
    --input_format=tf_frozen_model \
    --output_format=tfjs_graph_model \
    --quantization_bytes=1 \
    --inputs=input_node \
    --outputs=output_node \
    path/to/frozen_model.pb \
    path/to/tfjs_model
  • --inputs и --outputs — имена входных и выходных узлов графа.
  • --quantization_bytes — опциональный параметр для уменьшения размера весов (1, 2 или 4 байта).

Опции конвертера

  • --quantization_bytes — сжатие весов:

    • 1 — наименьший размер, но теряется точность.
    • 2 — компромисс между размером и точностью.
    • 4 — без сжатия, полная точность.
  • --skip_op_check — пропуск проверки поддерживаемых операций, полезно для моделей с нестандартными слоями.

  • --strip_debug_ops — удаление отладочных операций из графа, снижая размер модели.

  • --weight_shard_size_bytes — размер одного весового файла при шардировании.


Примеры использования в браузере

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

import * as tf from '@tensorflow/tfjs';

// Загрузка модели
const model = await tf.loadLayersModel('path/to/tfjs_model/model.json');

// Использование модели для предсказаний
const input = tf.tensor2d([[1, 2, 3, 4]]);
const output = model.predict(input);
output.print();

Для моделей графового формата (GraphModel) используется tf.loadGraphModel.


Советы по конвертации

  • Перед конвертацией проверить совместимость всех слоёв с TensorFlow.js. Не все пользовательские слои Keras поддерживаются напрямую.
  • Для больших моделей использовать --quantization_bytes=2 или 1 для уменьшения размера.
  • Проверять корректность работы модели после конвертации: иногда требуется подправить входные данные или имена узлов.
  • Хранить весовые файлы и model.json в одной директории для корректного доступа из JavaScript.

Особенности SavedModel с несколькими сигнатурами

Если SavedModel содержит несколько вариантов инференса, можно указать конкретный с помощью --saved_model_tags:

tensorflowjs_converter \
    --input_format=tf_saved_model \
    --saved_model_tags=serve,experimental \
    path/to/saved_model \
    path/to/tfjs_model

Это позволяет выбирать только необходимый граф для веб-приложений.


Работа с пользовательскими слоями Keras

При наличии кастомных слоёв:

  1. Добавить их в JavaScript через tf.serialization.registerClass.
  2. При загрузке модели указать класс:
tf.serialization.registerClass(CustomLayer);
const model = await tf.loadLayersModel('path/to/tfjs_model/model.json');

Без регистрации пользовательский слой вызовет ошибку при загрузке модели.


Вывод

Использование tensorflowjs_converter обеспечивает плавный переход от Python-моделей к веб-приложениям на JavaScript. Правильный выбор формата исходной модели, настройка тегов, параметров сжатия и учёт кастомных слоёв гарантируют стабильную работу моделей в среде TensorFlow.js.