Verstehen Sie die Op-Registrierung und Kernel-Verknüpfung in TensorFlow

Question

Ich bin ziemlich neu in TensorFlow und jetzt in die benutzerdefinierte OP-Entwicklung. Ich habe bereits das offizielle Tutorial gelesen, aber ich fühle, dass hinter den Kulissen viele Dinge passieren und ich möchte meine benutzerdefinierten Ops nicht immer in user_ops Verzeichnis einfügen.

Als solches ich nahm ein example word2vec

, die einen eigenen "Skipgram" op verwendet, deren Registrierung hier definiert ist:
/word2vec_ops.cc
und Kernel-Implementierung, die hier:
/word2vec_kernels.cc

Mit Blick auf die Build-Datei, ich habe versucht, individuelle Ziele

012 zu bauen

1) bazel build -c opt tensorflow/models/embedding:word2vec_ops
Dies generiert Bündel von Objektdateien wie erwartet.

2) bazel build -c opt tensorflow/models/embedding:word2vec_kernels
Gleiches für diese.

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

Dieser letzte Build verwendet eine benutzerdefinierte Regel nämlich tf_op_gen_wrapper_py https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorflow.bzl#L197-L231

Interessant zu beachten, dass dies nur auf Op Registrierung abhängig und nicht auf dem Kernel selbst.

Schließlich oben, wenn ich py_binary bauen sich

mit

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

es funktioniert gut, aber ich sehe nicht, wo und wie der Kernel C++ Code verknüpft?

Zusätzlich würde ich auch gerne die tf_op_gen_wrapper_py Regel und die ganze Compilation/Linking-Prozedur, die hinter den Kulissen für die Ops-Registrierung geht, verstehen.

Danke.

Answer 1

Wenn adding a new kind of operation to TensorFlow gibt es zwei Hauptschritte:

1. Registering the "op", die eine Schnittstelle für den Betrieb beinhaltet definieren, und

2. Registering one or more "kernels", die Umsetzung (en) für das umfasst die Definition Operation, vielleicht mit spezialisierten Implementierungen für verschiedene Datentypen oder Gerätetypen (wie CPU oder GPU).

Beide Schritte umfassen das Schreiben von C++ - Code. Das Registrieren eines Ops verwendet die REGISTER_OP() macro, und die Registrierung eines Kernels verwendet die REGISTER_KERNEL_BUILDER() macro. Diese Makros erstellen statische Initialisierer, die ausgeführt werden, wenn das Modul, das sie enthält, geladen wird.Es gibt zwei Hauptmechanismen für die Op- und Kernel-Registrierung:

Statische Verknüpfung in die Core-TensorFlow-Bibliothek und statische Initialisierung.

Dynamische Verknüpfung zur Laufzeit mit der Funktion tf.load_op_library().

Im Fall "Skipgram" verwenden wir Option 1 (statische Verknüpfung). Die Ops sind in die zentrale TensorFlow-Bibliothek here eingebunden, und die Kernel sind in here verknüpft. (Beachten Sie, dass dies nicht ideal ist: Die word2vec Ops wurden erstellt, bevor wir tf.load_op_library() hatten, und daher gab es keinen Mechanismus, um sie dynamisch zu verknüpfen.) Daher werden die Ops und Kernel beim ersten Laden von TensorFlow registriert (in import tensorflow as tf). Wenn sie heute erstellt würden, würden sie dynamisch geladen werden, so dass sie nur registriert würden, wenn sie benötigt würden. (Der Code hat eine SyntaxNet example dynamischer Belastung.)

Die tf_op_gen_wrapper_py() rule in Bazel nimmt eine Liste von op -Bibliothek Abhängigkeiten und erzeugt Python Wrapper für die ops. Der Grund, dass diese Regel nur von der Op-Registrierung abhängt, ist, dass die Python-Wrapper vollständig von der Schnittstelle des Operators bestimmt werden, der in der Op-Registrierung definiert ist. Insbesondere hat die Python-Schnittstelle keine Ahnung, ob es spezialisierte Kernel für einen bestimmten Typ oder ein bestimmtes Gerät gibt. Der Wrapper-Generator verknüpft die Op-Registrierungen in eine simple C++ binary, die Python-Code für jeden der registrierten Ops generiert. Beachten Sie, dass Sie, wenn Sie tf.load_op_library() verwenden, den Wrappergenerator nicht selbst aufrufen müssen, da tf.load_op_library() den erforderlichen Code zur Laufzeit generiert.