Zapoznaj się z rejestracją i łączeniem jądra w TensorFlow

Question

Jestem całkiem nowy w TensorFlow i właśnie teraz zajmuję się niestandardowym rozwojem op. Przeczytałem już oficjalny tutorial, ale czuję, że wiele rzeczy dzieje się za kulisami i nie zawsze chcę umieścić moje niestandardowe ops w katalogu user_ops.

Jako taki, podjąłem się example word2vec

który wykorzystuje niestandardowe "Skipgram" OP, którego rejestracja jest zdefiniowany tutaj:
/word2vec_ops.cc
i którego jądro realizacja jest tutaj:
/word2vec_kernels.cc

Patrząc na plik budowy, próbowałem zbudować pojedyncze cele

1) bazel build -c opt tensorflow/models/embedding:word2vec_ops
Generuje to wiązkę plików obiektów zgodnie z oczekiwaniami.

2) bazel build -c opt tensorflow/models/embedding:word2vec_kernels
To samo dotyczy tego.

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

Ten ostatni build używa niestandardową regułę mianowicie tf_op_gen_wrapper_py https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorflow.bzl#L197-L231

warto zauważyć, że to zależy tylko od op rejestracji, a nie na samym jądrze.

Wszakże powyżej, jeżeli buduję sam py_binary użyciu

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

to działa dobrze, ale nie widzę, gdzie i jak C++ kod jądra związane?

Dodatkowo, chciałbym również zrozumieć regułę tf_op_gen_wrapper_py i całą procedurę kompilacji/łączenia, która idzie w parze z rejestracją operacji.

Dzięki.

Answer 1

Gdy adding a new kind of operation to TensorFlow, istnieją dwa główne etapy:

1. Registering the "op", który obejmuje definiowanie interfejs operacji i

2. Registering one or more "kernels", który obejmuje określenie realizacji (-y) operacja, być może ze specjalistycznymi implementacjami dla różnych typów danych lub typów urządzeń (takich jak procesor lub GPU).

Oba etapy dotyczą pisania kodu C++. Zarejestrowanie operacji używa numeru REGISTER_OP() macro, a rejestracja jądra odbywa się przy użyciu REGISTER_KERNEL_BUILDER() macro. Makra te tworzą statyczne inicjalizatory, które działają, gdy załadowany jest moduł zawierający te moduły.Istnieją dwa główne mechanizmy op i jądra rejestracji:

1. Static łączące w podstawowej bibliotece TensorFlow i inicjalizacji statycznej.

2. Dynamiczne łączenie w czasie wykonywania przy użyciu funkcji tf.load_op_library().

W przypadku "Skipgram" używamy opcji 1 (łączenie statyczne). Operacje te są połączone z rdzeniową biblioteką TensorFlow here, a jądra są połączone w here. (Zauważ, że nie jest to idealne: ops word2vec zostały utworzone zanim mieliśmy tf.load_op_library(), więc nie było mechanizmu dynamicznego ich łączenia.) W związku z tym operacje i jądra są rejestrowane po pierwszym załadowaniu TensorFlow (w import tensorflow as tf). Gdyby zostały utworzone dzisiaj, zostałyby załadowane dynamicznie, tak, że byłyby zarejestrowane, gdyby były potrzebne. (Kod SyntaxNet ma example dynamicznego obciążenia.)

tf_op_gen_wrapper_py() rule w Bazel pobiera listę op -biblioteka zależnościami i generuje obwolut Python dla tych OPS. Powodem, że ta reguła zależy tylko od rejestracji op jest to, że owijacze Pythona są określone w całości przez interfejs operacji op, który jest zdefiniowany w op rejestracji. Warto zauważyć, że interfejs Pythona nie ma pojęcia, czy istnieją wyspecjalizowane jądra dla określonego typu lub urządzenia. Generator opakowujący łączy rejestracje operacji w simple C++ binary, która generuje kod Pythona dla każdej zarejestrowanej operacji. Pamiętaj, że jeśli używasz tf.load_op_library(), nie musisz samodzielnie wywoływać generatora opakowania, ponieważ tf.load_op_library() wygeneruje potrzebny kod w czasie wykonywania.