Понимание регистрации операций и связывания ядра в TensorFlow

Question

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

В качестве примера я взял пример word2vec

который использует пользовательский оп "Skipgram", регистрация которого определена здесь:
/word2vec_ops.cc
и чья реализация ядра здесь:
/word2vec_kernels.cc

Глядя на файл сборки, я пытался создать отдельные цели

1) bazel build -c opt tensorflow/models/embedding:word2vec_ops
Это создает кучу объектных файлов, как и ожидалось.

2) bazel build -c opt tensorflow/models/embedding:word2vec_kernels
То же самое для этого.

3) bazel build -c opt tensorflow/models/embedding:word2vec_kernels:gen_word2vec

Эта последняя сборка использует пользовательское правило, а именно tf_op_gen_wrapper_py https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorflow.bzl#L197-L231

Интересно отметить, что это зависит только от регистрации Op, а не от самого ядра.

Ведь выше, если я построю py_binary сам с помощью

bazel build -c opt tensorflow/models/embedding:word2vec

это работает нормально, но я не вижу, где и как связан код ядра C++?

Кроме того, я также хотел бы понять tf_op_gen_wrapper_py правило и вся процедура компиляции / компоновки, которая идет за кулисы для регистрации ops.

Благодарю.

Answer 1

При добавлении нового типа операции в TensorFlow, есть два основных шага:

1. Регистрация "операции", которая включает в себя определение интерфейса для операции, и

2. Регистрация одного или нескольких "ядер", которая включает в себя определение реализации (й) для операции, возможно, с помощью специализированных реализаций для различных типов данных или типов устройств (таких как CPU или GPU).

Оба этапа включают в себя написание кода на C++. Регистрация оп использует REGISTER_OP() макрос, и регистрация ядра использует REGISTER_KERNEL_BUILDER() макрос. Эти макросы создают статические инициализаторы, которые запускаются при загрузке модуля, содержащего их. Существует два основных механизма регистрации оп и ядра:

1. Статическое связывание с базовой библиотекой TensorFlow и статическая инициализация.

2. Динамическое связывание во время выполнения, используя tf.load_op_library() функция.

В случае "Skipgram" Мы используем вариант 1 (статическое связывание). Здесь операции связаны с основной библиотекой TensorFlow, а ядра связаны здесь. (Обратите внимание, что это не идеально: word2vec Операции были созданы до того, как мы tf.load_op_library() и поэтому не было никакого механизма для их динамического связывания.) Следовательно, ops и ядра регистрируются при первой загрузке TensorFlow (в import tensorflow as tf). Если бы они были созданы сегодня, они были бы динамически загружены, так что они были бы зарегистрированы только в случае необходимости. (В коде SyntaxNet есть пример динамической загрузки.)

tf_op_gen_wrapper_py() rule в Bazel берет список зависимостей op -library и генерирует обертки Python для этих операций. Причина, по которой это правило зависит только от регистрации операции, заключается в том, что оболочки Python полностью определяются интерфейсом операции, который определяется в регистрации операции. Примечательно, что интерфейс Python не знает, существуют ли специализированные ядра для определенного типа или устройства. Генератор оболочки связывает регистрацию операций в простой двоичный файл C++, который генерирует код Python для каждого из зарегистрированных операций. Обратите внимание, что если вы используете tf.load_op_library() Вам не нужно вызывать генератор-обертку самостоятельно, потому что tf.load_op_library() сгенерирует необходимый код во время выполнения.