Pisanie niestandardowych operacji, jąder i gradientów w TensorFlow.js

Przegląd

W tym przewodniku opisano mechanizmy definiowania niestandardowych operacji (ops), jąder i gradientów w TensorFlow.js. Ma na celu przedstawienie przeglądu głównych koncepcji i wskazówek do kodu, który demonstruje koncepcje w działaniu.

Dla kogo jest ten przewodnik?

Jest to dość zaawansowany przewodnik, który porusza pewne elementy wewnętrzne TensorFlow.js. Może być szczególnie przydatny dla następujących grup osób:

  • Zaawansowani użytkownicy TensorFlow.js zainteresowani dostosowywaniem zachowania różnych operacji matematycznych (np. badacze nadpisujący istniejące implementacje gradientów lub użytkownicy, którzy muszą załatać brakującą funkcjonalność w bibliotece)
  • Użytkownicy budujący biblioteki rozszerzające TensorFlow.js (np. ogólna biblioteka algebry liniowej zbudowana na elementach pierwotnych TensorFlow.js lub nowy backend TensorFlow.js).
  • Użytkownicy zainteresowani wnoszeniem nowych operacji do tensorflow.js, którzy chcą uzyskać ogólny przegląd działania tych mechanizmów.

To nie jest przewodnik po ogólnym użyciu TensorFlow.js, ponieważ dotyczy wewnętrznych mechanizmów wdrażania. Nie musisz rozumieć tych mechanizmów, aby korzystać z TensorFlow.js

Aby w pełni wykorzystać ten przewodnik, musisz znać się na czytaniu kodu źródłowego TensorFlow.js (lub chcieć spróbować).

Terminologia

W tym przewodniku warto opisać kilka kluczowych terminów.

Operacje (Ops) — operacja matematyczna na jednym lub większej liczbie tensorów, która daje jako wynik jeden lub więcej tensorów. Operacje to kod „wysokiego poziomu” i mogą używać innych operacji do definiowania swojej logiki.

Jądro — konkretna implementacja dostosowana do możliwości konkretnego sprzętu/platformy. Jądra są „niskiego poziomu” i specyficzne dla backendu. Niektóre operacje mają mapowanie jeden do jednego z op na jądro, podczas gdy inne korzystają z wielu jąder.

Gradient / GradFunc — definicja operacji/jądra w „trybie wstecz”, która oblicza pochodną tej funkcji w odniesieniu do niektórych danych wejściowych. Gradienty to kod „wysokiego poziomu” (nie specyficzny dla backendu) i mogą wywoływać inne operacje lub jądra.

Rejestr jądra — mapa krotki (nazwa jądra, nazwa zaplecza) na implementację jądra.

Rejestr gradientowy — mapa nazwy jądra na implementację gradientu .

Organizacja kodu

Operacje i gradienty są zdefiniowane w tfjs-core .

Jądra są specyficzne dla backendu i są zdefiniowane w odpowiednich folderach backendu (np. tfjs-backend-cpu ).

Niestandardowe operacje, jądra i gradienty nie muszą być definiowane w tych pakietach. Ale często będą używać podobnych symboli w swojej implementacji.

Wdrażanie niestandardowych operacji

Operację niestandardową można sobie wyobrazić na przykład jako funkcję JavaScript, która zwraca dane wyjściowe tensora, często z tensorami jako danymi wejściowymi.

  • Niektóre operacje można całkowicie zdefiniować w kategoriach istniejących operacji i należy po prostu bezpośrednio importować i wywoływać te funkcje. Oto przykład .
  • Implementacja operacji może również zostać wysłana do określonych jąder zaplecza. Odbywa się to poprzez Engine.runKernel i zostanie opisane dalej w sekcji „Implementowanie niestandardowych jąder”. Oto przykład .

Implementowanie niestandardowych jąder

Implementacje jądra specyficzne dla backendu pozwalają na zoptymalizowaną implementację logiki dla danej operacji. Jądra są wywoływane przez ops wywołujące tf.engine().runKernel() . Implementacje jądra definiują cztery rzeczy

  • Nazwa jądra.
  • Backend, w którym zaimplementowane jest jądro.
  • Wejścia: Argumenty tensorowe funkcji jądra.
  • Atrybuty: Argumenty niebędące tensorem funkcji jądra.

Oto przykład implementacji jądra . Konwencje użyte do wdrożenia są specyficzne dla backendu i najlepiej je zrozumieć, przyglądając się implementacji i dokumentacji każdego konkretnego backendu.

Ogólnie jądra działają na poziomie niższym niż tensory i zamiast tego bezpośrednio odczytują i zapisują w pamięci, która ostatecznie zostanie opakowana w tensory przez rdzeń tfjs.

Po zaimplementowaniu jądra można je zarejestrować w TensorFlow.js za pomocą funkcji registerKernel z tfjs-core. Możesz zarejestrować jądro dla każdego backendu, w którym ma działać. Po zarejestrowaniu jądro można wywołać za pomocą tf.engine().runKernel(...) a TensorFlow.js upewni się, że zostanie wysłany do implementacji w bieżący aktywny backend.

Implementowanie niestandardowych gradientów

Gradienty są ogólnie definiowane dla danego jądra (identyfikowane przez tę samą nazwę jądra, która jest używana w wywołaniu tf.engine().runKernel(...) ). Dzięki temu tfjs-core może używać rejestru do wyszukiwania definicji gradientów dla dowolnego jądra w czasie wykonywania.

Implementowanie niestandardowych gradientów jest przydatne w przypadku:

  • Dodanie definicji gradientu, której może nie być w bibliotece
  • Zastąpienie istniejącej definicji gradientu w celu dostosowania obliczeń gradientu dla danego jądra.

Przykłady implementacji gradientów można zobaczyć tutaj .

Po zaimplementowaniu gradientu dla danego wywołania można go zarejestrować w TensorFlow.js za pomocą funkcji registerGradient z tfjs-core.

Innym podejściem do implementowania niestandardowych gradientów, które omija rejestr gradientów (i w ten sposób pozwala na obliczanie gradientów dla dowolnych funkcji w dowolny sposób, jest użycie tf.customGrad .

Oto przykład operacji w bibliotece wykorzystującej niestandardowe Grad