Написание пользовательских операций, ядер и градиентов в TensorFlow.js, Написание пользовательских операций, ядер и градиентов в TensorFlow.js

Обзор

В этом руководстве описываются механизмы определения пользовательских операций (ops), ядер и градиентов в TensorFlow.js. Он призван предоставить обзор основных концепций и указателей на код, демонстрирующий концепции в действии.

Для кого это руководство?

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

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

Это не руководство по общему использованию TensorFlow.js, поскольку оно касается внутренних механизмов реализации. Вам не нужно понимать эти механизмы, чтобы использовать TensorFlow.js.

Вам нужно уметь (или хотеть попробовать) читать исходный код TensorFlow.js, чтобы максимально использовать это руководство.

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

Для этого руководства полезно описать несколько ключевых терминов заранее.

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

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

Gradient / GradFunc — определение «обратного режима» оператора /ядра , которое вычисляет производную этой функции относительно некоторых входных данных. Градиенты представляют собой код «высокого уровня» (не зависящий от бэкэнда) и могут вызывать другие операции или ядра.

Реестр ядра — сопоставление кортежа (имя ядра, имя серверной части) с реализацией ядра.

Gradient Registry — сопоставление имени ядра с реализацией градиента .

Организация кода

Операции и градиенты определены в 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.

,

Обзор

В этом руководстве описываются механизмы определения пользовательских операций (ops), ядер и градиентов в TensorFlow.js. Он призван предоставить обзор основных концепций и указателей на код, демонстрирующий концепции в действии.

Для кого это руководство?

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

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

Это не руководство по общему использованию TensorFlow.js, поскольку оно касается внутренних механизмов реализации. Вам не нужно понимать эти механизмы, чтобы использовать TensorFlow.js.

Вам нужно уметь (или хотеть попробовать) читать исходный код TensorFlow.js, чтобы максимально использовать это руководство.

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

Для этого руководства полезно описать несколько ключевых терминов заранее.

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

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

Gradient / GradFunc — определение «обратного режима» оператора /ядра , которое вычисляет производную этой функции относительно некоторых входных данных. Градиенты представляют собой код «высокого уровня» (не зависящий от бэкэнда) и могут вызывать другие операции или ядра.

Реестр ядра — сопоставление кортежа (имя ядра, имя серверной части) с реализацией ядра.

Gradient Registry — сопоставление имени ядра с реализацией градиента .

Организация кода

Операции и градиенты определены в 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.