Utwórz op

Jeśli chcesz utworzyć operację, która nie jest objęta istniejącą biblioteką TensorFlow, zalecamy najpierw spróbować napisać operację w Pythonie jako kompozycję istniejących operacji lub funkcji Pythona. Jeśli nie jest to możliwe, możesz utworzyć niestandardową operację C++. Istnieje kilka powodów, dla których warto utworzyć niestandardową operację C++:

  • Nie jest łatwo ani możliwe wyrazić swoją operację jako kompozycję istniejących operacji.
  • Wyrażanie operacji jako kompozycji istniejących prymitywów nie jest efektywne.
  • Chcesz ręcznie połączyć kompozycję prymitywów, która dla przyszłego kompilatora byłaby trudna.

Załóżmy na przykład, że chcesz zaimplementować coś w rodzaju „łączenia median”, podobnego do operatora „MaxPool”, ale obliczania median w przesuwających się oknach zamiast wartości maksymalnych. Wykonanie tego przy użyciu kombinacji operacji może być możliwe (np. przy użyciu ExtractImagePatches i TopK), ale może nie być tak wydajne i wydajne pod względem pamięci jak operacja natywna, w której można zrobić coś sprytniejszego w pojedynczej, połączonej operacji. Jak zawsze, zazwyczaj warto najpierw spróbować wyrazić to, co chcesz, używając kompozycji operatorów, a dodać nową operację tylko wtedy, gdy okaże się to trudne lub nieefektywne.

Aby uwzględnić niestandardową operację, musisz:

  1. Zarejestruj nową operację w pliku C++. Rejestracja operacji definiuje interfejs (specyfikację) funkcjonalności operacji, która jest niezależna od implementacji operacji. Na przykład rejestracja operacji definiuje nazwę operacji oraz dane wejściowe i wyjściowe operacji. Definiuje również funkcję kształtu używaną do wnioskowania o kształcie tensora.
  2. Zaimplementuj op w C++. Implementacja operacji jest nazywana jądrem i jest konkretną implementacją specyfikacji zarejestrowanej w kroku 1. Może istnieć wiele jąder dla różnych typów wejść/wyjść lub architektur (na przykład procesorów CPU, GPU).
  3. Utwórz opakowanie Pythona (opcjonalnie). To opakowanie to publiczny interfejs API używany do tworzenia operacji w Pythonie. Na podstawie rejestracji op generowane jest domyślne opakowanie, którego można użyć bezpośrednio lub dodać do niego.
  4. Napisz funkcję obliczającą gradienty dla operacji (opcjonalnie).
  5. Przetestuj op. Zwykle robimy to w Pythonie dla wygody, ale możesz także przetestować tę operację w C++. Jeśli zdefiniujesz gradienty, możesz je zweryfikować za pomocą Pythona tf.test.compute_gradient_error . Zobacz relu_op_test.py jako przykład testujący funkcje forward operatorów podobnych do Relu i ich gradienty.

Warunki wstępne

Zdefiniuj interfejs operacyjny

Definiujesz interfejs operacji rejestrując ją w systemie TensorFlow. Podczas rejestracji określasz nazwę swojej operacji, jej dane wejściowe (typy i nazwy) i dane wyjściowe (typy i nazwy), a także dokumenty i wszelkie atrybuty, których operacja może wymagać.

Aby zobaczyć, jak to działa, załóżmy, że chcesz utworzyć operację, która pobiera tensor int32 s i wyprowadza kopię tensora, ze wszystkimi elementami oprócz pierwszego ustawionymi na zero. W tym celu utwórz plik o nazwie zero_out.cc . Następnie dodaj wywołanie do makra REGISTER_OP , które definiuje interfejs dla Twojej operacji:

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

Ta operacja ZeroOut przyjmuje jako dane wejściowe jeden tensor to_zero 32-bitowych liczb całkowitych i wyprowadza tensor zeroed dla 32-bitowych liczb całkowitych. Operacja wykorzystuje również funkcję kształtu, aby upewnić się, że tensor wyjściowy ma taki sam kształt jak tensor wejściowy. Na przykład, jeśli danymi wejściowymi jest tensor kształtu [10, 20], wówczas ta funkcja kształtu określa, że ​​kształt wyjściowy również ma postać [10, 20].

Zaimplementuj jądro dla op

Po zdefiniowaniu interfejsu podaj jedną lub więcej implementacji op. Aby utworzyć jedno z tych jąder, utwórz klasę rozszerzającą OpKernel i zastępując metodę Compute . Metoda Compute udostępnia jeden argument context typu OpKernelContext* , z którego można uzyskać dostęp do przydatnych rzeczy, takich jak tensory wejściowe i wyjściowe.

Dodaj jądro do pliku, który utworzyłeś powyżej. Jądro może wyglądać mniej więcej tak:

#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<int32>();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output_flat = output_tensor->flat<int32>();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    for (int i = 1; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value if possible.
    if (N > 0) output_flat(0) = input(0);
  }
};

Po zaimplementowaniu swojego jądra rejestrujesz je w systemie TensorFlow. Podczas rejestracji określasz różne ograniczenia, zgodnie z którymi będzie działać to jądro. Na przykład możesz mieć jedno jądro przeznaczone dla procesorów i osobne dla procesorów graficznych.

Aby to zrobić dla operacji ZeroOut , dodaj następujące polecenie do zero_out.cc :

REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

Wielowątkowe jądra procesora

Aby napisać wielowątkowe jądro procesora, można użyć funkcji Shard w work_sharder.h . Ta funkcja dzieli funkcję obliczeniową na wątki skonfigurowane do użycia na potrzeby wątków wewnątrz operacji (zobacz intra_op_parallelism_threads w config.proto ).

Jądra GPU

Jądro GPU składa się z dwóch części: OpKernel i jądra CUDA oraz jego kodu uruchamiającego.

Czasami implementacja OpKernel jest wspólna dla jądra procesora i procesora graficznego, na przykład podczas sprawdzania danych wejściowych i przydzielania wyników. W takim przypadku sugerowaną implementacją jest:

  1. Zdefiniuj szablon OpKernel na urządzeniu i pierwotny typ tensora.
  2. Aby dokonać faktycznego obliczenia wyniku, funkcja Compute wywołuje szablonową strukturę funktora.
  3. Specjalizacja tego funktora dla CPUDevice jest zdefiniowana w tym samym pliku, ale specjalizacja dla GPUDevice jest zdefiniowana w pliku .cu.cc, ponieważ zostanie ona skompilowana za pomocą kompilatora CUDA.

Oto przykładowa realizacja. 

// kernel_example.h
#ifndef KERNEL_EXAMPLE_H_
#define KERNEL_EXAMPLE_H_

#include <unsupported/Eigen/CXX11/Tensor>

template <typename Device, typename T>
struct ExampleFunctor {
  void operator()(const Device& d, int size, const T* in, T* out);
};

#if GOOGLE_CUDA
// Partially specialize functor for GpuDevice.
template <typename T>
struct ExampleFunctor<Eigen::GpuDevice, T> {
  void operator()(const Eigen::GpuDevice& d, int size, const T* in, T* out);
};
#endif

#endif KERNEL_EXAMPLE_H_

// kernel_example.cc
#include "kernel_example.h"

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;

REGISTER_OP("Example")
    .Attr("T: numbertype")
    .Input("input: T")
    .Output("input_times_two: T")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

// CPU specialization of actual computation.
template <typename T>
struct ExampleFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, int size, const T* in, T* out) {
    for (int i = 0; i < size; ++i) {
      out[i] = 2 * in[i];
    }
  }
};

// OpKernel definition.
// template parameter <T> is the datatype of the tensors.
template <typename Device, typename T>
class ExampleOp : public OpKernel {
 public:
  explicit ExampleOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));

    // Do the computation.
    OP_REQUIRES(context, input_tensor.NumElements() <= tensorflow::kint32max,
                errors::InvalidArgument("Too many elements in tensor"));
    ExampleFunctor<Device, T>()(
        context->eigen_device<Device>(),
        static_cast<int>(input_tensor.NumElements()),
        input_tensor.flat<T>().data(),
        output_tensor->flat<T>().data());
  }
};

// Register the CPU kernels.
#define REGISTER_CPU(T)                                          \
  REGISTER_KERNEL_BUILDER(                                       \
      Name("Example").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
      ExampleOp<CPUDevice, T>);
REGISTER_CPU(float);
REGISTER_CPU(int32);

// Register the GPU kernels.
#ifdef GOOGLE_CUDA
#define REGISTER_GPU(T)                                          \
  /* Declare explicit instantiations in kernel_example.cu.cc. */ \
  extern template class ExampleFunctor<GPUDevice, T>;            \
  REGISTER_KERNEL_BUILDER(                                       \
      Name("Example").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
      ExampleOp<GPUDevice, T>);
REGISTER_GPU(float);
REGISTER_GPU(int32);
#endif  // GOOGLE_CUDA

// kernel_example.cu.cc
#ifdef GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "kernel_example.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"

using namespace tensorflow;

using GPUDevice = Eigen::GpuDevice;

// Define the CUDA kernel.
template <typename T>
__global__ void ExampleCudaKernel(const int size, const T* in, T* out) {
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size;
       i += blockDim.x * gridDim.x) {
    out[i] = 2 * __ldg(in + i);
  }
}

// Define the GPU implementation that launches the CUDA kernel.
template <typename T>
void ExampleFunctor<GPUDevice, T>::operator()(
    const GPUDevice& d, int size, const T* in, T* out) {
  // Launch the cuda kernel.
  //
  // See core/util/gpu_kernel_helper.h for example of computing
  // block count and thread_per_block count.
  int block_count = 1024;
  int thread_per_block = 20;
  ExampleCudaKernel<T>
      <<<block_count, thread_per_block, 0, d.stream()>>>(size, in, out);
}

// Explicitly instantiate functors for the types of OpKernels registered.
template struct ExampleFunctor<GPUDevice, float>;
template struct ExampleFunctor<GPUDevice, int32>;

#endif  // GOOGLE_CUDA

Zbuduj bibliotekę op

Skompiluj op za pomocą kompilatora systemu (instalacja binarna TensorFlow)

Powinieneś być w stanie skompilować zero_out.cc za pomocą kompilatora C++ takiego jak g++ lub clang dostępnego w twoim systemie. Binarny pakiet PIP instaluje pliki nagłówkowe i bibliotekę potrzebną do skompilowania operacji w lokalizacjach specyficznych dla systemu. Jednakże biblioteka Pythona TensorFlow udostępnia funkcję get_include umożliwiającą pobranie katalogu nagłówkowego, a katalog get_lib zawiera obiekt współdzielony, z którym można się połączyć. Oto wyniki tych funkcji na komputerze Ubuntu.

$ python
>>> import tensorflow as tf
>>> tf.sysconfig.get_include()
'/usr/local/lib/python3.6/site-packages/tensorflow/include'
>>> tf.sysconfig.get_lib()
'/usr/local/lib/python3.6/site-packages/tensorflow'

Zakładając, że masz zainstalowany g++ , oto sekwencja poleceń, których możesz użyć do skompilowania operacji w bibliotekę dynamiczną.

TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )
g++ -std=c++14 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

W systemie macOS podczas tworzenia pliku .so wymagana jest dodatkowa flaga „-undefinied dynamic_lookup”.

Uwaga na temat wersji gcc >=5 : gcc używa nowego ABI C++ od wersji 5 . TensorFlow 2.8 i wcześniejsze zostały zbudowane przy użyciu gcc4 , które korzysta ze starszego ABI. Jeśli używasz tych wersji TensorFlow i próbujesz skompilować bibliotekę op za pomocą gcc>=5 , dodaj -D_GLIBCXX_USE_CXX11_ABI=0 do wiersza poleceń, aby biblioteka była kompatybilna ze starszym ABI. Pakiety TensorFlow 2.9+ są domyślnie kompatybilne z nowszym ABI.

Skompiluj op za pomocą bazela (instalacja źródłowa TensorFlow)

Jeśli masz zainstalowane źródła TensorFlow, możesz skorzystać z systemu kompilacji TensorFlow, aby skompilować swój op. Umieść plik BUILD z następującą regułą kompilacji Bazel w katalogu tensorflow/core/user_ops .

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    name = "zero_out.so",
    srcs = ["zero_out.cc"],
)

Uruchom następujące polecenie, aby zbudować zero_out.so .

$ bazel build --config opt //tensorflow/core/user_ops:zero_out.so

Aby skompilować operację Example z jądrem CUDA, należy użyć parametru gpu_srcs z tf_custom_op_library . Umieść plik BUILD z następującą regułą kompilacji Bazel w nowym folderze w katalogu tensorflow/core/user_ops (np. „example_gpu”).

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    # kernel_example.cc  kernel_example.cu.cc  kernel_example.h
    name = "kernel_example.so",
    srcs = ["kernel_example.h", "kernel_example.cc"],
    gpu_srcs = ["kernel_example.cu.cc", "kernel_example.h"],
)

Uruchom następującą komendę, aby skompilować kernel_example.so .

$ bazel build --config opt //tensorflow/core/user_ops/example_gpu:kernel_example.so

Użyj op w Pythonie

TensorFlow Python API udostępnia funkcję tf.load_op_library służącą do ładowania biblioteki dynamicznej i rejestrowania operacji w środowisku TensorFlow. load_op_library zwraca moduł Pythona, który zawiera opakowania Pythona dla operacji i jądra. Zatem po zbudowaniu operacji możesz wykonać następujące czynności, aby uruchomić ją z poziomu Pythona:

import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')
print(zero_out_module.zero_out([[1, 2], [3, 4]]).numpy())

# Prints
array([[1, 0], [0, 0]], dtype=int32)

Należy pamiętać, że wygenerowana funkcja otrzyma nazwę Snake_case (aby zachować zgodność z PEP8 ). Tak więc, jeśli twoja operacja ma nazwę ZeroOut w plikach C++, funkcja Pythona będzie nazywana zero_out .

Aby udostępnić tę opcję jako zwykłą funkcję import z modułu Pythona, przydatne może być wywołanie funkcji load_op_library w pliku źródłowym Pythona w następujący sposób: 

import tensorflow as tf

zero_out_module = tf.load_op_library('./zero_out.so')
zero_out = zero_out_module.zero_out

Sprawdź, czy operacja działa

Dobrym sposobem sprawdzenia, czy pomyślnie zaimplementowałeś operację, jest napisanie dla niej testu. Utwórz plik zero_out_op_test.py o zawartości:

import tensorflow as tf

class ZeroOutTest(tf.test.TestCase):
  def testZeroOut(self):
    zero_out_module = tf.load_op_library('./zero_out.so')
    with self.test_session():
      result = zero_out_module.zero_out([5, 4, 3, 2, 1])
      self.assertAllEqual(result.eval(), [5, 0, 0, 0, 0])

if __name__ == "__main__":
  tf.test.main()

Następnie uruchom test (zakładając, że masz zainstalowany tensorflow): 

$ python zero_out_op_test.py

Wbuduj zaawansowane funkcje w swój op

Teraz, gdy wiesz, jak zbudować podstawową (i nieco ograniczoną) operację i jej implementację, przyjrzymy się niektórym z bardziej skomplikowanych rzeczy, które zazwyczaj będziesz musiał wbudować w swoją operację. Obejmuje to:

Kontrole warunkowe i walidacja

W powyższym przykładzie założono, że op ma zastosowanie do tensora o dowolnym kształcie. A co by było, gdyby dotyczyło to tylko wektorów? Oznacza to dodanie kontroli do powyższej implementacji OpKernel.

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(input_tensor.shape()),
                errors::InvalidArgument("ZeroOut expects a 1-D vector."));
    // ...
  }

To potwierdza, że ​​dane wejściowe są wektorem i zwraca po ustawieniu statusu InvalidArgument , jeśli tak nie jest. Makro OP_REQUIRES przyjmuje trzy argumenty:

Alternatywnie, jeśli chcesz sprawdzić, czy obiekt Status zwrócony przez jakąś funkcję jest błędem, a jeśli tak, zwróć go, użyj OP_REQUIRES_OK . Obydwa te makra wracają z funkcji w przypadku błędu.

Rejestracja op

Atrybuty

Operacje mogą mieć atrybuty, których wartości są ustawiane podczas dodawania operacji do wykresu. Służą one do konfigurowania operacji, a dostęp do ich wartości można uzyskać zarówno w implementacji jądra, jak i w typach wejść i wyjść w rejestracji operacji. Jeśli to możliwe, preferuj używanie danych wejściowych zamiast attr, ponieważ dane wejściowe są bardziej elastyczne. Dzieje się tak, ponieważ atrybuty są stałymi i muszą zostać zdefiniowane w momencie tworzenia wykresu. Natomiast dane wejściowe są tensorami, których wartości mogą być dynamiczne; oznacza to, że dane wejściowe mogą zmieniać się w każdym kroku, być ustawiane przy użyciu źródła danych itp. Atrybuty są używane do rzeczy, których nie można zrobić za pomocą wejść: dowolna konfiguracja, która wpływa na sygnaturę (liczba lub typ wejść lub wyjść) lub która może' nie zmieniać się krok po kroku.

Definiujesz atrybut podczas rejestracji operacji, określając jego nazwę i typ za pomocą metody Attr , która oczekuje specyfikacji w postaci:

<name>: <attr-type-expr>

gdzie <name> zaczyna się od litery i może składać się ze znaków alfanumerycznych i podkreśleń, a <attr-type-expr> jest wyrażeniem typu w postaci opisanej poniżej .

Na przykład, jeśli chcesz, aby operacja ZeroOut zachowywała indeks określony przez użytkownika, zamiast tylko elementu zerowego, możesz zarejestrować tę operację w następujący sposób:

REGISTER_OP("ZeroOut")
    .Attr("preserve_index: int")
    .Input("to_zero: int32")
    .Output("zeroed: int32");

(Zauważ, że zestaw typów atrybutów różni się od tf.DType używanego dla wejść i wyjść.)

Twoje jądro może następnie uzyskać dostęp do tego atrybutu w swoim konstruktorze poprzez parametr context :

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {
    // Get the index of the value to preserve
    OP_REQUIRES_OK(context,
                   context->GetAttr("preserve_index", &preserve_index_));
    // Check that preserve_index is positive
    OP_REQUIRES(context, preserve_index_ >= 0,
                errors::InvalidArgument("Need preserve_index >= 0, got ",
                                        preserve_index_));
  }
  void Compute(OpKernelContext* context) override {
    // ...
  }
 private:
  int preserve_index_;
};

które można następnie wykorzystać w metodzie Compute : 

  void Compute(OpKernelContext* context) override {
    // ...

    // We're using saved attr to validate potentially dynamic input
    // So we check that preserve_index is in range
    OP_REQUIRES(context, preserve_index_ < input.dimension(0),
                errors::InvalidArgument("preserve_index out of range"));

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the requested input value
    output_flat(preserve_index_) = input(preserve_index_);
  }

Typy atrybutów

W attr obsługiwane są następujące typy:

  • string : Dowolna sekwencja bajtów (nie musi to być UTF8).
  • int : Liczba całkowita ze znakiem.
  • float : Liczba zmiennoprzecinkowa.
  • bool : Prawda czy fałsz.
  • type : Jedna z (innych niż ref) wartości DataType .
  • shape : TensorShapeProto .
  • list(<type>) : Lista <type> , gdzie <type> jest jednym z powyższych typów. Zauważ, że list(list(<type>)) jest niepoprawna.

Zobacz także: op_def_builder.cc:FinalizeAttr , aby uzyskać ostateczną listę.

Wartości domyślne i ograniczenia

Atrybuty mogą mieć wartości domyślne, a niektóre typy atrybutów mogą mieć ograniczenia. Aby zdefiniować atrybut z ograniczeniami, możesz użyć następujących <attr-type-expr> s:

{'<string1>', '<string2>'} : Wartość musi być ciągiem znaków o wartości <string1> lub <string2> . W przypadku użycia tej składni implikowana jest nazwa typu string . To emuluje wyliczenie:

REGISTER_OP("EnumExample")
    .Attr("e: {'apple', 'orange'}");

{<type1>, <type2>} : wartość jest typu type i musi być jedną z <type1> lub <type2> , gdzie <type1> i <type2> są obsługiwane tf.DType . Nie określasz, że typem attr jest type . Jest to domniemane, gdy masz listę typów w {...} . Na przykład w tym przypadku attr t jest typem, który musi być int32 , float lub bool :

REGISTER_OP("RestrictedTypeExample")
    .Attr("t: {int32, float, bool}");

Istnieją skróty do typowych ograniczeń typu:

  • numbertype : type typu ograniczony do typów numerycznych (nieciągowych i boolowych).
  • realnumbertype : Podobnie jak numbertype bez typów złożonych.
  • quantizedtype : Podobnie jak numbertype , ale tylko typy liczbowe skwantowane.

Konkretne listy typów dozwolonych przez te typy są definiowane przez funkcje (takie jak NumberTypes() ) w tensorflow/core/framework/types.h . W tym przykładzie atrybut t musi należeć do jednego z typów liczbowych:

REGISTER_OP("NumberType")
    .Attr("t: numbertype");

Dla tej operacji:

tf.number_type(t=tf.int32)  # Valid
tf.number_type(t=tf.bool)   # Invalid

Listy można łączyć z innymi listami i pojedynczymi typami. Poniższa operacja pozwala, aby attr t był dowolnym typem numerycznym lub typem bool:

REGISTER_OP("NumberOrBooleanType")
    .Attr("t: {numbertype, bool}");

Dla tej operacji:

tf.number_or_boolean_type(t=tf.int32)  # Valid
tf.number_or_boolean_type(t=tf.bool)   # Valid
tf.number_or_boolean_type(t=tf.string) # Invalid

int >= <n> : Wartość musi być int, której wartość jest większa lub równa <n> , gdzie <n> jest liczbą naturalną. Na przykład poniższa rejestracja op określa, że ​​atrybut a musi mieć wartość co najmniej 2 :

REGISTER_OP("MinIntExample")
    .Attr("a: int >= 2");

list(<type>) >= <n> : Lista typu <type> , której długość jest większa lub równa <n> . Na przykład poniższa rejestracja op określa, że ​​attr a jest listą typów ( int32 lub float ) i że muszą być ich co najmniej 3:

REGISTER_OP("TypeListExample")
    .Attr("a: list({int32, float}) >= 3");

Aby ustawić domyślną wartość atrybutu (czyniąc go opcjonalnym w wygenerowanym kodzie), dodaj na końcu = <default> , jak w:

REGISTER_OP("AttrDefaultExample")
    .Attr("i: int = 0");

Dodatkowo można określić zarówno ograniczenie, jak i wartość domyślną:

REGISTER_OP("AttrConstraintAndDefaultExample")
    .Attr("i: int >= 1 = 1");

Obsługiwana składnia wartości domyślnej jest używana w proto reprezentacji wynikowej definicji GraphDef.

Oto przykłady określenia wartości domyślnej dla wszystkich typów:

REGISTER_OP("AttrDefaultExampleForAllTypes")
   .Attr("s: string = 'foo'")
   .Attr("i: int = 0")
   .Attr("f: float = 1.0")
   .Attr("b: bool = true")
   .Attr("ty: type = DT_INT32")
   .Attr("sh: shape = { dim { size: 1 } dim { size: 2 } }")
   .Attr("te: tensor = { dtype: DT_INT32 int_val: 5 }")
   .Attr("l_empty: list(int) = []")
   .Attr("l_int: list(int) = [2, 3, 5, 7]");

Należy w szczególności zauważyć, że wartości typu type używają tf.DType .

Wielopostaciowość

Typ polimorfizmu

W przypadku operacji, które mogą przyjmować różne typy danych wejściowych lub generować różne typy wyników, można określić atrybut w typie danych wejściowych lub wyjściowych w rejestracji operacji. Zwykle rejestruje się wówczas OpKernel dla każdego obsługiwanego typu.

Na przykład, jeśli chcesz, aby operacja ZeroOut działała na float oprócz int32 , rejestracja Twojej operacji może wyglądać następująco:

REGISTER_OP("ZeroOut")
    .Attr("T: {float, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

Twoja rejestracja op określa teraz, że typem danych wejściowych musi być float lub int32 i że jego dane wyjściowe będą tego samego typu, ponieważ oba mają typ T

Nazewnictwo

Wejściom, wyjściom i atrybutom należy ogólnie nadać nazwy Snake_case. Jedynym wyjątkiem są atrybuty używane jako typ wejścia lub typ wyjścia. Te atrybuty można wywnioskować po dodaniu operacji do wykresu, dlatego nie pojawiają się w funkcji operacji. Na przykład ta ostatnia definicja ZeroOut wygeneruje funkcję Pythona, która wygląda następująco:

def zero_out(to_zero, name=None):
  """...
  Args:
    to_zero: A `Tensor`. Must be one of the following types:
        `float32`, `int32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `to_zero`.
  """

Jeśli to_zero zostanie przekazany tensor int32 , wówczas T zostanie automatycznie ustawione na int32 (no cóż, właściwie DT_INT32 ). Tym wywnioskowanym atrybutom nadawane są nazwy pisane wielką literą lub CamelCase.

Porównaj to z op, która ma atrybut typu, który określa typ wyjściowy:

REGISTER_OP("StringToNumber")
    .Input("string_tensor: string")
    .Output("output: out_type")
    .Attr("out_type: {float, int32} = DT_FLOAT");
    .Doc(R"doc(
Converts each string in the input Tensor to the specified numeric type.
)doc");

W tym przypadku użytkownik musi określić typ wyjścia, tak jak w wygenerowanym Pythonie: 

def string_to_number(string_tensor, out_type=None, name=None):
  """Converts each string in the input Tensor to the specified numeric type.

  Args:
    string_tensor: A `Tensor` of type `string`.
    out_type: An optional `tf.DType` from: `tf.float32, tf.int32`.
      Defaults to `tf.float32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor` of type `out_type`.
  """
Wpisz przykład polimorfizmu
#include "tensorflow/core/framework/op_kernel.h"

class ZeroOutInt32Op : public OpKernel {
  // as before
};

class ZeroOutFloatOp : public OpKernel {
 public:
  explicit ZeroOutFloatOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<float>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<float>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint<int32>("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<int32>("T"),
    ZeroOutInt32Op);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutFloatOp);

Aby zachować kompatybilność wsteczną , powinieneś określić wartość domyślną podczas dodawania atrybutu do istniejącej operacji:

REGISTER_OP("ZeroOut")
  .Attr("T: {float, int32} = DT_INT32")
  .Input("to_zero: T")
  .Output("zeroed: T")

Powiedzmy, że chcesz dodać więcej typów, powiedzmy double :

REGISTER_OP("ZeroOut")
    .Attr("T: {float, double, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

Zamiast pisać kolejny OpKernel ze zbędnym kodem, jak powyżej, często będziesz mógł zamiast tego użyć szablonu C++. Nadal będziesz mieć jedną rejestrację jądra (wywołanie REGISTER_KERNEL_BUILDER ) na każde przeciążenie.

template <typename T>
class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<T>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<T>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint<int32>("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<int32>("T"),
    ZeroOutOp<int32>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<double>("T"),
    ZeroOutOp<double>);

Jeśli masz więcej niż kilka przeciążeń, możesz umieścić rejestrację w makrze.

#include "tensorflow/core/framework/op_kernel.h"

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

REGISTER_KERNEL(int32);
REGISTER_KERNEL(float);
REGISTER_KERNEL(double);

#undef REGISTER_KERNEL

W zależności od listy typów, dla których rejestrujesz jądro, możesz użyć makra dostarczonego przez tensorflow/core/framework/register_types.h : 

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"

REGISTER_OP("ZeroOut")
    .Attr("T: realnumbertype")
    .Input("to_zero: T")
    .Output("zeroed: T");

template <typename T>
class ZeroOutOp : public OpKernel { ... };

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);

#undef REGISTER_KERNEL
Lista wejść i wyjść

Oprócz możliwości akceptowania lub tworzenia różnych typów, ops mogą zużywać lub produkować zmienną liczbę tensorów.

W następnym przykładzie attr T przechowuje listę typów i jest używany jako typ in i out . Dane wejściowe i wyjściowe są listami tensorów tego typu (a liczba i typy tensorów na wyjściu są takie same jak na wejściu, ponieważ oba mają typ T ).

REGISTER_OP("PolymorphicListExample")
    .Attr("T: list(type)")
    .Input("in: T")
    .Output("out: T");

Możesz także nałożyć ograniczenia na typy, które można określić na liście. W następnym przypadku danymi wejściowymi jest lista tensorów float i double . Operacja akceptuje na przykład typy danych wejściowych (float, double, float) i w takim przypadku typem wyjściowym będzie również (float, double, float) .

REGISTER_OP("ListTypeRestrictionExample")
    .Attr("T: list({float, double})")
    .Input("in: T")
    .Output("out: T");

Jeśli chcesz, aby wszystkie tensory na liście były tego samego typu, możesz zrobić coś takiego:

REGISTER_OP("IntListInputExample")
    .Attr("N: int")
    .Input("in: N * int32")
    .Output("out: int32");

Akceptuje listę tensorów int32 i używa int attr N do określenia długości listy.

Można to również uczynić typem polimorficznym . W następnym przykładzie danymi wejściowymi jest lista tensorów (o długości "N" ) tego samego (ale nieokreślonego) typu ( "T" ), a danymi wyjściowymi jest pojedynczy tensor pasującego typu:

REGISTER_OP("SameListInputExample")
    .Attr("N: int")
    .Attr("T: type")
    .Input("in: N * T")
    .Output("out: T");

Domyślnie listy tensorów mają minimalną długość 1. Możesz zmienić to ustawienie domyślne, używając ograniczenia ">=" na odpowiednim attr . W następnym przykładzie danymi wejściowymi jest lista co najmniej 2 tensorów int32 :

REGISTER_OP("MinLengthIntListExample")
    .Attr("N: int >= 2")
    .Input("in: N * int32")
    .Output("out: int32");

Ta sama składnia działa z atrybutami "list(type)" : 

REGISTER_OP("MinimumLengthPolymorphicListExample")
    .Attr("T: list(type) >= 3")
    .Input("in: T")
    .Output("out: T");

Wejścia i wyjścia

Podsumowując powyższe, rejestracja op może mieć wiele danych wejściowych i wyjściowych:

REGISTER_OP("MultipleInsAndOuts")
    .Input("y: int32")
    .Input("z: float")
    .Output("a: string")
    .Output("b: int32");

Każda specyfikacja wejścia lub wyjścia ma postać:

<name>: <io-type-expr>

gdzie <name> zaczyna się od litery i może składać się ze znaków alfanumerycznych i podkreśleń. <io-type-expr> jest jednym z następujących wyrażeń typu:

  • <type> , gdzie <type> jest obsługiwanym typem danych wejściowych (np. float , int32 , string ). Określa pojedynczy tensor danego typu.

    Zobacz tf.DType .

    REGISTER_OP("BuiltInTypesExample")
    .Input("integers: int32")
    .Input("complex_numbers: complex64");

  • <attr-type> , gdzie <attr-type> to nazwa atrybutu typu type lub list(type) (z możliwym ograniczeniem typu). Ta składnia pozwala na operacje polimorficzne .

    REGISTER_OP("PolymorphicSingleInput")
    .Attr("T: type")
    .Input("in: T");

REGISTER_OP("RestrictedPolymorphicSingleInput")
    .Attr("T: {int32, int64}")
    .Input("in: T");

    Odwoływanie się do atrybutu typu list(type) pozwala zaakceptować sekwencję tensorów.

    REGISTER_OP("ArbitraryTensorSequenceExample")
    .Attr("T: list(type)")
    .Input("in: T")
    .Output("out: T");

REGISTER_OP("RestrictedTensorSequenceExample")
    .Attr("T: list({int32, int64})")
    .Input("in: T")
    .Output("out: T");

    Należy zauważyć, że liczba i typy tensorów na wyjściu out są takie same jak na wejściu in , ponieważ oba są typu T .

  • Dla sekwencji tensorów tego samego typu: <number> * <type> , gdzie <number> to nazwa atrybutu typu int . <type> może być albo tf.DType , albo nazwą atrybutu o typie type . Jako przykład pierwszej, ta operacja akceptuje listę tensorów int32 :

    REGISTER_OP("Int32SequenceExample")
    .Attr("NumTensors: int")
    .Input("in: NumTensors * int32")

    Podczas gdy ta opcja akceptuje listę tensorów dowolnego typu, o ile wszystkie są takie same:

    REGISTER_OP("SameTypeSequenceExample")
    .Attr("NumTensors: int")
    .Attr("T: type")
    .Input("in: NumTensors * T")

  • Aby uzyskać odwołanie do tensora: Ref(<type>) , gdzie <type> jest jednym z poprzednich typów.

Każdy atrybut użyty w typie wejścia zostanie wywnioskowany. Zgodnie z konwencją te wywnioskowane atrybuty używają wielkich nazw (takich jak T lub N ). W przeciwnym razie wejścia, wyjścia i atrybuty mają nazwy podobne do parametrów funkcji (np. num_outputs ). Aby uzyskać więcej informacji, zobacz wcześniejszą sekcję dotyczącą nazewnictwa .

Aby uzyskać więcej informacji, zobacz tensorflow/core/framework/op_def_builder.h .

Kompatybilność wsteczna

Załóżmy, że napisałeś ładną, niestandardową operację i udostępniłeś ją innym, dzięki czemu masz zadowolonych klientów korzystających z Twojej operacji. Jednakże chciałbyś w jakiś sposób wprowadzić zmiany w operacji.

Ogólnie rzecz biorąc, zmiany w istniejących, sprawdzonych specyfikacjach muszą być kompatybilne wstecz: zmiana specyfikacji operacji nie może przerwać wcześniejszych serializowanych buforów protokołu GraphDef zbudowanych na podstawie starszych specyfikacji. Szczegóły kompatybilności GraphDef są opisane tutaj .

Istnieje kilka sposobów zachowania kompatybilności wstecznej.

  1. Wszelkie nowe atrybuty dodane do operacji muszą mieć zdefiniowane wartości domyślne, a przy tej wartości domyślnej operacja musi zachować oryginalne zachowanie. Aby zmienić operację z niepolimorficznej na polimorficzną, musisz nadać wartość domyślną nowemu atrybutowi typu, aby domyślnie zachować oryginalny podpis. Na przykład, jeśli Twoja operacja była:

    REGISTER_OP("MyGeneralUnaryOp")
    .Input("in: float")
    .Output("out: float");

    możesz uczynić go polimorficznym w sposób kompatybilny wstecz, używając:

    REGISTER_OP("MyGeneralUnaryOp")
    .Input("in: T")
    .Output("out: T")
    .Attr("T: numerictype = DT_FLOAT");

  2. Możesz bezpiecznie ograniczyć atrybut, mniej restrykcyjny. Na przykład możesz zmienić z {int32, int64} na {int32, int64, float} lub type . Możesz też zmienić z {"apple", "orange"} na {"apple", "banana", "orange"} lub string .

  3. Możesz zamienić pojedyncze wejścia/wyjścia na listę wejść/wyjść, o ile domyślny typ listy odpowiada staremu podpisowi.

  4. Możesz dodać nowe wejście/wyjście listy, jeśli domyślnie jest puste.

  5. Przestrzeń nazw dla wszystkich nowych operacji, które tworzysz, poprzedzając nazwy operacji czymś unikalnym dla twojego projektu. Pozwala to uniknąć kolizji operacji z operacjami, które mogą zostać uwzględnione w przyszłych wersjach TensorFlow.

  6. Planuj z wyprzedzeniem! Spróbuj przewidzieć przyszłe zastosowania op. Niektórych zmian w podpisie nie można przeprowadzić w zgodny sposób (na przykład przekształcenie listy tego samego typu w listę różnych typów).

Pełną listę bezpiecznych i niebezpiecznych zmian można znaleźć w tensorflow/core/framework/op_compatibility_test.cc . Jeśli nie możesz zapewnić zgodności zmiany operacji z poprzednimi wersjami, utwórz nową operację z nową nazwą i nową semantyką.

Należy również pamiętać, że chociaż te zmiany mogą zachować zgodność GraphDef , wygenerowany kod Pythona może zmienić się w sposób, który nie jest kompatybilny ze starymi programami wywołującymi. Interfejs API języka Python można zachować zgodność poprzez ostrożne zmiany w odręcznie napisanym opakowaniu języka Python, zachowując stary podpis, z wyjątkiem ewentualnego dodania na końcu nowych opcjonalnych argumentów. Ogólnie niezgodne zmiany można wprowadzać tylko wtedy, gdy TensorFlow zmienia główne wersje i muszą być zgodne z semantyką wersji GraphDef .

Wsparcie GPU

Możesz zaimplementować różne OpKernels i zarejestrować jeden dla procesora, a drugi dla GPU, tak samo jak możesz zarejestrować jądra dla różnych typów . Istnieje kilka przykładów jąder z obsługą GPU w tensorflow/core/kernels/ . Zwróć uwagę, że niektóre jądra mają wersję procesora w pliku .cc , wersję procesora graficznego w pliku kończącym się na _gpu.cu.cc i część wspólnego kodu w pliku .h .

Na przykład tf.pad ma wszystko oprócz jądra GPU w tensorflow/core/kernels/pad_op.cc . Jądro GPU znajduje się w tensorflow/core/kernels/pad_op_gpu.cu.cc , a udostępniony kod to klasa z szablonem zdefiniowana w tensorflow/core/kernels/pad_op.h . Organizujemy kod w ten sposób z dwóch powodów: pozwala to na współdzielenie wspólnego kodu pomiędzy implementacjami CPU i GPU oraz umieszcza implementację GPU w oddzielnym pliku, dzięki czemu może być skompilowana tylko przez kompilator GPU.

Należy zauważyć, że nawet jeśli używana jest wersja pad z jądrem GPU, nadal potrzebuje on wejścia "paddings" do pamięci procesora. Aby oznaczyć, że wejścia lub wyjścia są przechowywane w CPU, dodaj wywołanie HostMemory() do rejestracji jądra, np.: 

#define REGISTER_GPU_KERNEL(T)                         \
  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
                              .Device(DEVICE_GPU)      \
                              .TypeConstraint<T>("T")  \
                              .HostMemory("paddings"), \
                          PadOp<GPUDevice, T>)

Kompilacja jądra dla urządzenia GPU

Spójrz na cuda_op_kernel.cu.cc, aby zobaczyć przykład użycia jądra CUDA do implementacji op. tf_custom_op_library akceptuje argument gpu_srcs , w którym można określić listę plików źródłowych zawierających jądra CUDA (pliki *.cu.cc ). Aby można było używać z binarną instalacją TensorFlow, jądra CUDA muszą zostać skompilowane za pomocą kompilatora nvcc firmy NVIDIA. Oto sekwencja poleceń, których możesz użyć do skompilowania cuda_op_kernel.cu.cc i cuda_op_kernel.cc w jedną, dynamicznie ładowaną bibliotekę:

nvcc -std=c++14 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \
  ${TF_CFLAGS[@]} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC

g++ -std=c++14 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \
  cuda_op_kernel.cu.o ${TF_CFLAGS[@]} -fPIC -lcudart ${TF_LFLAGS[@]}

cuda_op_kernel.so utworzony powyżej można załadować w Pythonie w zwykły sposób, używając funkcji tf.load_op_library .

Zauważ, że jeśli biblioteki CUDA nie są zainstalowane w /usr/local/lib64 , musisz jawnie określić ścieżkę w drugim poleceniu (g++) powyżej. Na przykład dodaj -L /usr/local/cuda-8.0/lib64/ jeśli CUDA jest zainstalowana w /usr/local/cuda-8.0 .

Zaimplementuj gradient w Pythonie

Biorąc pod uwagę wykres operacji, TensorFlow wykorzystuje automatyczne różnicowanie (propagację wsteczną), aby dodać nowe operacje reprezentujące gradienty w stosunku do istniejących operacji. Aby automatyczne różnicowanie działało w przypadku nowych operacji, musisz zarejestrować funkcję gradientu, która oblicza gradienty w odniesieniu do danych wejściowych operacji, biorąc pod uwagę gradienty w odniesieniu do wyników operacji.

Matematycznie, jeśli op oblicza \(y = f(x)\) zarejestrowany gradient op konwertuje gradienty \(\partial L/ \partial y\) straty \(L\) w odniesieniu do\(y\) w gradienty \(\partial L/ \partial x\) w odniesieniu do \(x\) za pomocą reguły łańcuchowej:

\[\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial f}{\partial x}.\]

W przypadku ZeroOut tylko jeden wpis na wejściu wpływa na wynik, więc gradient w stosunku do wejścia jest rzadkim tensorem „jednego gorącego”. Wyraża się to w następujący sposób:

from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops

@ops.RegisterGradient("ZeroOut")
def _zero_out_grad(op, grad):
  """The gradients for `zero_out`.

  Args:
    op: The `zero_out` `Operation` that we are differentiating, which we can use
      to find the inputs and outputs of the original op.
    grad: Gradient with respect to the output of the `zero_out` op.

  Returns:
    Gradients with respect to the input of `zero_out`.
  """
  to_zero = op.inputs[0]
  shape = array_ops.shape(to_zero)
  index = array_ops.zeros_like(shape)
  first_grad = array_ops.reshape(grad, [-1])[0]
  to_zero_grad = sparse_ops.sparse_to_dense([index], shape, first_grad, 0)
  return [to_zero_grad]  # List of one Tensor, since we have one input

Szczegóły dotyczące rejestrowania funkcji gradientowych za pomocą tf.RegisterGradient :

  • W przypadku operacji z jednym wyjściem funkcja gradientu pobierze tf.Operation , op i tf.Tensor grad i zbuduje nowe operacje z tensorów op.inputs[i] , op.outputs[i] i grad . Informacje o dowolnych atrybutach można znaleźć poprzez tf.Operation.get_attr .

  • Jeśli operacja ma wiele wyników, funkcja gradientu przyjmie op i grads , gdzie grads to lista gradientów w odniesieniu do każdego wyjścia. Wynikiem funkcji gradientu musi być lista obiektów Tensor reprezentujących gradienty w odniesieniu do każdego wejścia.

  • Jeśli dla niektórych danych wejściowych nie ma dobrze zdefiniowanego gradientu, na przykład dla danych wejściowych w postaci liczb całkowitych używanych jako indeksy, odpowiadający zwrócony gradient powinien mieć None . Na przykład, dla operacji przyjmującej tensor zmiennoprzecinkowy x i indeks całkowity i , funkcja gradientu return [x_grad, None] .

  • Jeśli w ogóle nie ma znaczącego gradientu dla operacji, często nie będziesz musiał rejestrować żadnego gradientu i dopóki gradient operacji nie będzie nigdy potrzebny, wszystko będzie dobrze. W niektórych przypadkach operacja nie ma dobrze zdefiniowanego gradientu, ale może zostać wykorzystana w obliczeniu gradientu. Tutaj możesz użyć ops.NotDifferentiable do automatycznego propagowania zer wstecz.

Należy zauważyć, że w momencie wywołania funkcji gradientu dostępny jest tylko wykres przepływu danych operacji, a nie same dane tensora. Dlatego wszystkie obliczenia muszą być wykonywane przy użyciu innych operacji tensorflow, które mają być uruchamiane w czasie wykonywania wykresu.

Dodaj wskazówki dotyczące typu podczas rejestrowania niestandardowego gradientu dla typu operacji, aby uczynić kod bardziej czytelnym, debugowalnym, łatwiejszym w utrzymaniu i solidniejszym dzięki sprawdzaniu poprawności danych. Na przykład, jeśli przyjmujesz op jako parametr funkcji, określ, że funkcja gradientu przyjmie tf.Operation jako typ parametru.

Funkcje kształtu w C++

Interfejs API TensorFlow ma funkcję zwaną „wnioskowaniem o kształcie”, która dostarcza informacji o kształtach tensorów bez konieczności wykonywania wykresu. Wnioskowanie o kształcie jest obsługiwane przez „funkcje kształtu”, które są zarejestrowane dla każdego typu operacji w deklaracji C++ REGISTER_OP i pełnią dwie role: stwierdzają, że kształty danych wejściowych są kompatybilne podczas konstruowania grafu i określają kształty wyników.

Funkcje kształtu są zdefiniowane jako operacje na klasie shape_inference::InferenceContext . Na przykład w funkcji kształtu dla ZeroOut:

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

c->set_output(0, c->input(0)); deklaruje, że kształt pierwszego wyjścia powinien być ustawiony na kształt pierwszego wejścia. Jeżeli wyjście jest wybrane według indeksu jak w powyższym przykładzie, drugim parametrem set_output powinien być obiekt ShapeHandle . Możesz utworzyć pusty obiekt ShapeHandle za pomocą jego domyślnego konstruktora. Obiekt ShapeHandle dla wejścia o indeksie idx można uzyskać poprzez c->input(idx) .

Istnieje wiele typowych funkcji kształtu, które mają zastosowanie do wielu operacji, takich jak shape_inference::UnchangedShape , które można znaleźć w pliku common_shape_fns.h i których można używać w następujący sposób:

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);

Funkcja kształtu może również ograniczać kształt danych wejściowych. W przypadku wersji ZeroOut z ograniczeniem kształtu wektorowego funkcja kształtu wyglądałaby następująco:

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input));
      c->set_output(0, input);
      return Status::OK();
    });

Wywołanie WithRank sprawdza, czy kształt wejściowy c->input(0) ma kształt o dokładnie jednym wymiarze (lub jeśli kształt wejściowy jest nieznany, kształt wyjściowy będzie wektorem o jednym nieznanym wymiarze).

Jeśli Twoja operacja jest polimorficzna z wieloma danymi wejściowymi , możesz użyć elementów InferenceContext do określenia liczby kształtów do sprawdzenia, a Merge do sprawdzenia, czy wszystkie kształty są kompatybilne (alternatywnie możesz uzyskać dostęp do atrybutów wskazujących długości za pomocą InferenceContext::GetAttr , który zapewnia dostęp do atrybutów op).

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      ::tensorflow::shape_inference::ShapeHandle output;
      for (size_t i = 0; i < c->num_inputs(); ++i) {
        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 2, &input));
        TF_RETURN_IF_ERROR(c->Merge(output, input, &output));
      }
      c->set_output(0, output);
      return Status::OK();
    });

Ponieważ wnioskowanie o kształcie jest funkcją opcjonalną, a kształty tensorów mogą zmieniać się dynamicznie, funkcje kształtu muszą być odporne na niekompletne informacje o kształcie dla dowolnego z danych wejściowych. Metoda Merge w InferenceContext umożliwia wywołującemu stwierdzenie, że dwa kształty są takie same, nawet jeśli jeden lub oba nie mają pełnych informacji. Funkcje kształtu są zdefiniowane dla wszystkich podstawowych operacji TensorFlow i zapewniają wiele różnych przykładów użycia.

Klasa InferenceContext posiada wiele funkcji, których można użyć do zdefiniowania manipulacji funkcjami kształtu. Na przykład możesz sprawdzić, czy dany wymiar ma bardzo konkretną wartość, używając InferenceContext::Dim i InferenceContext::WithValue ; możesz określić, że wymiar wyjściowy jest sumą/iloczynem dwóch wymiarów wejściowych, używając InferenceContext::Add i InferenceContext::Multiply . Zobacz klasę InferenceContext aby poznać wszystkie różne manipulacje kształtami, które możesz określić. Poniższy przykład ustawia kształt pierwszego wyjścia na (n, 3), gdzie pierwsze wejście ma kształt (n, ...)

.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    c->set_output(0, c->Matrix(c->Dim(c->input(0), 0), 3));
    return Status::OK();
});

Jeśli masz skomplikowaną funkcję kształtu, powinieneś rozważyć dodanie testu sprawdzającego, czy różne kombinacje kształtów wejściowych dają oczekiwane kombinacje kształtów wyjściowych. Przykłady pisania tych testów można znaleźć w niektórych naszych testach podstawowych operacji . (Składnia INFER_OK i INFER_ERROR jest nieco zagadkowa, ale staraj się zachować zwięzłość w przedstawianiu specyfikacji kształtu wejściowego i wyjściowego w testach. Na razie przejrzyj komentarze wokół tych testów, aby uzyskać pojęcie o specyfikacji ciągu kształtu).

Zbuduj pakiet pip dla swojej niestandardowej operacji

Aby zbudować pakiet pip dla swojej operacji, zobacz przykład tensorflow/custom-op . W tym przewodniku pokazano, jak tworzyć niestandardowe operacje z pakietu pip TensorFlow zamiast kompilować TensorFlow ze źródła.