Обзор
В этом руководстве описываются механизмы определения пользовательских операций (ops), ядер и градиентов в TensorFlow.js. Его цель — предоставить обзор основных концепций и указатели на код, демонстрирующие эти концепции в действии.
Для кого предназначено это руководство?
Это довольно продвинутое руководство, затрагивающее некоторые внутренние особенности TensorFlow.js. Оно может быть особенно полезно для следующих групп людей:
- Опытные пользователи TensorFlow.js, заинтересованные в настройке поведения различных математических операций (например, исследователи, переопределяющие существующие реализации градиента, или пользователи, которым необходимо исправить недостающие функции в библиотеке).
- Пользователи, создающие библиотеки, расширяющие TensorFlow.js (например, общую библиотеку линейной алгебры, построенную на основе примитивов TensorFlow.js, или новый бэкэнд TensorFlow.js).
- Пользователи, заинтересованные в добавлении новых операций в tensorflow.js и желающие получить общее представление о том, как работают эти механизмы.
Это не руководство по общему использованию TensorFlow.js, поскольку оно посвящено внутренним механизмам реализации. Вам не нужно понимать эти механизмы, чтобы использовать TensorFlow.js.
Вам нужно уметь (или хотеть попробовать) читать исходный код TensorFlow.js, чтобы максимально эффективно использовать это руководство.
Терминология
В этом руководстве полезно заранее описать несколько ключевых терминов.
Операции (Ops) — математическая операция над одним или несколькими тензорами, которая создает на выходе один или несколько тензоров. Операции представляют собой код «высокого уровня» и могут использовать другие операции для определения своей логики.
Ядро — конкретная реализация операции, привязанная к конкретным возможностям оборудования/платформы. Ядра являются «низкоуровневыми» и зависят от серверной части. Некоторые операции имеют однозначное сопоставление операций с ядром, в то время как другие операции используют несколько ядер.
Gradient /GradFunc — определение операции/ядра в «обратном режиме», которое вычисляет производную этой функции относительно некоторого входного сигнала. Градиенты представляют собой код «высокого уровня» (не зависящий от серверной части) и могут вызывать другие операции или ядра.
Реестр ядра — сопоставление кортежа (имя ядра, имя бэкенда) с реализацией ядра.
Реестр градиентов — сопоставление имени ядра с реализацией градиента .
Организация кода
Операции и градиенты определены в tfjs-core .
Ядра зависят от серверной части и определяются в соответствующих папках серверной части (например, tfjs-backend-cpu ).
Внутри этих пакетов не нужно определять пользовательские операции, ядра и градиенты. Но в своей реализации часто будут использовать похожие символы.
Реализация пользовательских операций
Пользовательскую операцию можно рассматривать как функцию JavaScript, которая возвращает некоторый тензорный вывод, часто с тензорами в качестве входных данных.
- Некоторые операции могут быть полностью определены в терминах существующих операций, и их следует просто импортировать и вызывать эти функции напрямую. Вот пример .
- Реализация операции также может отправляться на определенные ядра бэкэнда. Это делается через
Engine.runKernelи будет описано далее в разделе «Реализация пользовательских ядер». Вот пример .
Реализация пользовательских ядер
Реализации ядра, специфичные для серверной части, позволяют оптимизировать реализацию логики для данной операции. Ядра вызываются операторами, вызывающими tf.engine().runKernel() . Реализация ядра определяется четырьмя вещами
- Имя ядра.
- Бэкэнд, в котором реализовано ядро.
- Входные данные: тензорные аргументы функции ядра.
- Атрибуты: нетензорные аргументы функции ядра.
Вот пример реализации ядра . Соглашения, используемые для реализации, зависят от конкретной серверной части и лучше всего понятны при рассмотрении реализации и документации каждой конкретной серверной части.
Обычно ядра работают на уровне ниже, чем тензоры, и вместо этого напрямую читают и записывают в память, которая в конечном итоге будет преобразована в тензоры с помощью tfjs-core.
После реализации ядра его можно зарегистрировать в TensorFlow.js с помощью функции registerKernel из tfjs-core. Вы можете зарегистрировать ядро для каждого бэкэнда, в котором вы хотите, чтобы это ядро работало. После регистрации ядро можно вызвать с помощью tf.engine().runKernel(...) и TensorFlow.js обязательно отправит его реализации в текущий активный бэкэнд.
Реализация пользовательских градиентов
Градиенты обычно определяются для данного ядра (идентифицируются тем же именем ядра, которое используется при вызове tf.engine().runKernel(...) ). Это позволяет tfjs-core использовать реестр для поиска определений градиентов для любого ядра во время выполнения.
Реализация пользовательских градиентов полезна для:
- Добавление определения градиента, которого может отсутствовать в библиотеке
- Переопределение существующего определения градиента для настройки вычисления градиента для данного ядра.
Посмотреть примеры реализации градиента можно здесь .
После того как вы реализовали градиент для данного вызова, его можно зарегистрировать в TensorFlow.js с помощью функции registerGradient из tfjs-core.
Другой подход к реализации пользовательских градиентов, который обходит реестр градиентов (и, таким образом, позволяет вычислять градиенты для произвольных функций произвольными способами, — это использование tf.customGrad .
Вот пример операции в библиотеке с использованием customGrad.