Создание оптимизированных по размеру пакетов браузера с помощью TensorFlow.js

Обзор

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

Эта функция предназначена для пользователей с производственными вариантами использования, которым будет особенно полезно сократить количество байтов в полезной нагрузке (и, следовательно, они готовы приложить усилия для достижения этой цели). Чтобы использовать эту функцию, вы должны быть знакомы с модулями ES , инструментами связывания JavaScript, такими как веб-пакет или накопительный пакет , а также такими концепциями, как устранение встряхивания деревьев/мертвого кода .

В этом руководстве показано, как создать собственный модуль tensorflow.js, который можно использовать с упаковщиком для создания оптимизированной по размеру сборки для программы с использованием tensorflow.js.

Терминология

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

Модули ESстандартная система модулей JavaScript . Представлено в ES6/ES2015. Идентифицируется по заявлениям об импорте и экспорте .

Объединение — получение набора ресурсов JavaScript и группировка/объединение их в один или несколько ресурсов JavaScript, которые можно использовать в браузере. На этом этапе обычно создаются конечные ресурсы, которые передаются браузеру. Приложения обычно создают свои собственные пакеты непосредственно из перекомпилированных источников библиотеки . К распространенным упаковщикам относятся накопительный пакет и веб-пакет . Конечным результатом объединения является пакет (или иногда фрагмент , если он разделен на несколько частей).

Удаление дерева/мертвого кода — удаление кода, который не используется в окончательном написанном приложении. Это делается во время объединения, обычно на этапе минификации.

Операции (Ops) — математическая операция над одним или несколькими тензорами, которая создает на выходе один или несколько тензоров. Операции представляют собой код «высокого уровня» и могут использовать другие операции для определения своей логики.

Ядро — конкретная реализация операции, привязанная к конкретным возможностям оборудования. Ядра являются «низкоуровневыми» и зависят от серверной части. Некоторые операции имеют однозначное сопоставление операций с ядром, в то время как другие операции используют несколько ядер.

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

Графические модели только для вывода

Основной вариант использования, о котором мы слышали от пользователей, связанных с этим, и который поддерживаем в этом выпуске, — это выполнение вывода с помощью графовых моделей TensorFlow.js . Если вы используете модель слоев TensorFlow.js , вы можете преобразовать ее в формат графовой модели с помощью tfjs-converter . Формат графовой модели более эффективен для варианта использования вывода.

Низкоуровневые манипуляции с тензором с помощью tfjs-core

Другой вариант использования, который мы поддерживаем, — это программы, которые напрямую используют пакет @tensorflow/tjfs-core для манипуляций с тензорами нижнего уровня.

Наш подход к индивидуальным конструкциям

Наши основные принципы при разработке этой функциональности включают следующее:

  • Максимально используйте систему модулей JavaScript (ESM) и позвольте пользователям TensorFlow.js делать то же самое.
  • Сделайте TensorFlow.js как можно более древовидным с помощью существующих сборщиков (например, веб-пакета, накопительного пакета и т. д.). Это позволяет пользователям воспользоваться всеми возможностями этих сборщиков, включая такие функции, как разделение кода.
  • Насколько это возможно, сохраняйте простоту использования для пользователей, которые не столь чувствительны к размеру пакета . Это означает, что производственные сборки потребуют больше усилий, поскольку многие настройки по умолчанию в наших библиотеках поддерживают простоту использования по сравнению со сборками, оптимизированными по размеру.

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

Хотя мы в первую очередь полагаемся на систему модулей JavaScript, мы также предоставляем собственный инструмент CLI для обработки частей, которые нелегко указать через систему модулей в коде, ориентированном на пользователя. Два примера:

  • Спецификации модели хранятся в файлах model.json .
  • Используемая нами система диспетчеризации для конкретного ядра.

Это делает создание пользовательской сборки tfjs немного более сложным, чем просто указание упаковщика на обычный пакет @tensorflow/tfjs.

Как создавать специальные пакеты с оптимизированным размером

Шаг 1. Определите, какие ядра использует ваша программа.

Этот шаг позволяет нам определить все ядра, используемые любыми запускаемыми вами моделями или кодом предварительной/постобработки, с учетом выбранного вами бэкэнда.

Используйте tf.profile для запуска частей вашего приложения, использующих tensorflow.js, и получения ядер. Это будет выглядеть примерно так

const profileInfo = await tf.profile(() => {
  // You must profile all uses of tf symbols.
  runAllMyTfjsCode();
});

const kernelNames = profileInfo.kernelNames
console.log(kernelNames);

Скопируйте этот список ядер в буфер обмена для следующего шага.

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

Вам нужно будет повторить этот шаг, если ваша модель изменится или изменится код предварительной/постобработки.

Шаг 2. Напишите файл конфигурации для пользовательского модуля tfjs.

Вот пример файла конфигурации.

Это выглядит так:

{
  "kernels": ["Reshape", "_FusedMatMul", "Identity"],
  "backends": [
      "cpu"
  ],
  "models": [
      "./model/model.json"
  ],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}
  • ядра: список ядер, которые будут включены в комплект. Скопируйте это из вывода шага 1.
  • бэкэнды: список бэкэндов, которые вы хотите включить. Допустимые параметры: «cpu», «webgl» и «wasm».
  • модели: список файлов model.json для моделей, которые вы загружаете в свое приложение. Может быть пустым, если ваша программа не использует tfjs_converter для загрузки графовой модели.
  • выходной путь: путь к папке для размещения сгенерированных модулей.
  • frontModeOnly: установите для этого параметра значение false, если вы хотите включить градиенты для ядер, перечисленных ранее.

Шаг 3. Создайте собственный модуль tfjs.

Запустите специальный инструмент сборки с файлом конфигурации в качестве аргумента. Чтобы иметь доступ к этому инструменту, вам необходимо установить пакет @tensorflow/tfjs .

npx tfjs-custom-module  --config custom_tfjs_config.json

Это создаст папку в outputPath с некоторыми новыми файлами.

Шаг 4. Настройте сборщик на псевдоним tfjs для нового пользовательского модуля.

В упаковщиках, таких как Webpack и Rollup, мы можем использовать псевдонимы существующих ссылок на модули tfjs, чтобы они указывали на наши недавно созданные пользовательские модули tfjs. Есть три модуля, которым необходимо связать псевдонимы для максимальной экономии размера пакета.

Вот фрагмент того, как это выглядит в веб-пакете ( полный пример здесь ):

...

config.resolve = {
  alias: {
    '@tensorflow/tfjs$':
        path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    '@tensorflow/tfjs-core$': path.resolve(
        __dirname, './custom_tfjs/custom_tfjs_core.js'),
    '@tensorflow/tfjs-core/dist/ops/ops_for_converter': path.resolve(
        __dirname, './custom_tfjs/custom_ops_for_converter.js'),
  }
}

...

А вот эквивалентный фрагмент кода для объединения ( полный пример здесь ):

import alias from '@rollup/plugin-alias';

...

alias({
  entries: [
    {
      find: /@tensorflow\/tfjs$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    },
    {
      find: /@tensorflow\/tfjs-core$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs_core.js'),
    },
    {
      find: '@tensorflow/tfjs-core/dist/ops/ops_for_converter',
      replacement: path.resolve(__dirname, './custom_tfjs/custom_ops_for_converter.js'),
    },
  ],
}));

...

Если ваш сборщик не поддерживает псевдонимы модулей, вам нужно будет изменить операторы import , чтобы импортировать tensorflow.js из созданного custom_tfjs.js , созданного на шаге 3. Определения операций не будут вытеснены в виде дерева, но ядра по-прежнему будут древовидными. - потрясен. Как правило, ядра, использующие древовидную тряску, обеспечивают наибольшую экономию конечного размера пакета.

Если вы используете только пакет @tensoflow/tfjs-core, вам нужно создать псевдоним только для этого одного пакета.

Шаг 5. Создайте свой пакет

Запустите сборщик (например webpack или rollup ), чтобы создать пакет. Размер пакета должен быть меньше, чем если бы вы запускали сборщик без псевдонимов модулей. Вы также можете использовать такие визуализаторы, чтобы увидеть, что попало в ваш окончательный пакет.

Шаг 6. Проверьте свое приложение

Обязательно проверьте, работает ли ваше приложение должным образом!