Utwórz op

Zadbaj o dobrą organizację dzięki kolekcji Zapisuj i kategoryzuj treści zgodnie ze swoimi preferencjami.

Jeśli chcesz utworzyć operację, której nie obejmuje istniejąca biblioteka 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 możesz chcieć stworzyć niestandardową opcję w C++:

  • Nie jest łatwo ani nie można wyrazić swojej operacji jako kompozycji istniejących operacji.
  • Wyrażanie swojej operacji jako kompozycji istniejących prymitywów nie jest efektywne.
  • Chcesz ręcznie połączyć kompozycję prymitywów, którą przyszły kompilator mógłby uznać za trudny do łączenia.

Na przykład wyobraź sobie, że chcesz zaimplementować coś takiego jak "łączenie median", podobne do operatora "MaxPool", ale obliczanie median w przesuwanych oknach zamiast wartości maksymalnych. Wykonanie tego przy użyciu kompozycji operacji może być możliwe (np. przy użyciu ExtractImagePatches i TopK), ale może nie być tak wydajne lub wydajne w zakresie pamięci, jak operacja natywna, w której można zrobić coś bardziej sprytnego w pojedynczej, połączonej operacji. Jak zawsze, zazwyczaj najpierw warto spróbować wyrazić to, co chcesz, używając składu operatora, wybierając tylko nową operację, jeśli okaże się to trudne lub nieefektywne.

Aby włączyć swoją niestandardową operację, musisz:

  1. Zarejestruj nową operację w pliku C++. Rejestracja operacji definiuje interfejs (specyfikację) dla funkcjonalności operacji, która jest niezależna od implementacji operacji. Na przykład rejestracja operacji definiuje nazwę operacji oraz wejścia i wyjścia operacji. Definiuje również funkcję kształtu używaną do wnioskowania o kształcie tensora.
  2. Zaimplementuj op w C++. Implementacja operacji nazywana jest jądrem i jest konkretną implementacją specyfikacji zarejestrowanej w kroku 1. Może istnieć wiele jąder dla różnych typów wejścia/wyjścia lub architektur (na przykład CPU, GPU).
  3. Utwórz opakowanie Pythona (opcjonalnie). Ten wrapper jest publicznym API używanym do tworzenia operacji w Pythonie. Z rejestracji operacji generowany jest domyślny wrapper, którego można użyć bezpośrednio lub dodać do niego.
  4. Napisz funkcję do obliczania gradientów dla operacji (opcjonalnie).
  5. Przetestuj op. Zwykle robimy to w Pythonie dla wygody, ale możesz też przetestować op 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, który testuje funkcje forward operatorów podobnych do Relu i ich gradienty.

Warunki wstępne

Zdefiniuj interfejs operacyjny

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

Aby zobaczyć, jak to działa, załóżmy, że chciałbyś utworzyć operację, która pobiera tensor z int32 s i wypisuje kopię tensora, przy czym wszystkie elementy oprócz pierwszego są ustawione na zero. Aby to zrobić, 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 ZeroOut pobiera jeden tensor to_zero 32-bitowych liczb całkowitych jako dane wejściowe i wyprowadza tensor zeroed z 32-bitowych liczb całkowitych. Op używa również funkcji kształtu, aby zapewnić, że tensor wyjściowy ma taki sam kształt jak tensor wejściowy. Na przykład, jeśli dane wejściowe są tensorem kształtu [10, 20], to ta funkcja kształtu określa, że ​​kształt wyjściowy to również [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ę, która rozszerza OpKernel i zastępuje 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 swoje 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 wdrożeniu jądra rejestrujesz je w systemie TensorFlow. Podczas rejestracji określasz różne ograniczenia, pod którymi to jądro będzie działać. Na przykład możesz mieć jedno jądro stworzone dla procesorów i osobne dla GPU.

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

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

Jądra procesora wielowątkowego

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 do wątków wewnątrzoperacyjnych (zobacz intra_op_parallelism_threads w config.proto ).

Jądra GPU

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

Czasami implementacja OpKernel jest wspólna dla jądra CPU i GPU, na przykład wokół sprawdzania wejść i przydzielania wyjść. W takim przypadku sugerowaną implementacją jest:

  1. Zdefiniuj szablon OpKernel na urządzeniu i typ pierwotny tensora.
  2. Aby wykonać rzeczywiste obliczenia danych wyjściowych, funkcja Compute wywołuje strukturę funktora z szablonem.
  3. Specjalizacja tego funktora dla CPUDevice jest zdefiniowana w tym samym pliku, ale specjalizacja dla GPUDevice jest zdefiniowana w pliku .cu.cc, ponieważ zostanie on skompilowany za pomocą kompilatora CUDA.

Oto przykładowa implementacja.

// 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 systemowego (instalacja binarna TensorFlow)

Powinieneś być w stanie skompilować zero_out.cc z kompilatorem C++ takim jak g++ lub clang dostępnym w twoim systemie. Pakiet binarny PIP instaluje pliki nagłówkowe i bibliotekę potrzebną do skompilowania operacji w lokalizacjach specyficznych dla systemu. Jednak biblioteka Pythona TensorFlow udostępnia funkcję get_include , aby uzyskać katalog nagłówka, a katalog get_lib ma obiekt współdzielony, z którym można się połączyć. Oto dane wyjściowe tych funkcji na komputerze z systemem 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 zainstalowane g++ , oto sekwencja poleceń, których możesz użyć do skompilowania swojej operacji do biblioteki dynamicznej.

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 „-undefined dynamic_lookup”.

Uwaga dotycząca wersji gcc >=5 : gcc używa nowego C++ ABI od wersji 5 . Pakiety binarne pip dostępne na stronie TensorFlow są zbudowane przy użyciu gcc4 , który używa starszego ABI. Jeśli kompilujesz bibliotekę op z gcc>=5 , dodaj -D_GLIBCXX_USE_CXX11_ABI=0 do wiersza poleceń, aby biblioteka była kompatybilna ze starszym abi.

Skompiluj op za pomocą bazel (instalacja źródła TensorFlow)

Jeśli masz zainstalowane źródła TensorFlow, możesz skorzystać z systemu kompilacji TensorFlow do skompilowania swojej operacji. Umieść plik BUILD z następującą regułą budowania 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, musisz 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ące polecenie, aby zbudować 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 do ładowania biblioteki dynamicznej i rejestrowania operacji we frameworku TensorFlow. load_op_library zwraca moduł Pythona, który zawiera opakowania Pythona dla operacji i jądra. Tak więc, kiedy już zbudujesz op, możesz wykonać następujące czynności, aby uruchomić go z 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)

Pamiętaj, że wygenerowana funkcja otrzyma nazwę snake_case (zgodną z PEP8 ). Tak więc, jeśli twoja operacja ma nazwę ZeroOut w plikach C++, funkcja Pythona będzie nazywać się zero_out .

Aby opcja była dostępna jako zwykły import funkcji z modułu Pythona, może być przydatne wywołanie 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 udało Ci się zaimplementować swoją operację, jest napisanie dla niej testu. Utwórz plik zero_out_op_test.py o treś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 swoją operację

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

Kontrole warunkowe i walidacja

W powyższym przykładzie założono, że op dotyczy tensora o dowolnym kształcie. Co by było, gdyby dotyczyło 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."));
    // ...
  }

Zapewnia to, że dane wejściowe są wektorem i zwraca po ustawieniu stanu InvalidArgument , jeśli tak nie jest. Makro OP_REQUIRES przyjmuje trzy argumenty:

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

Rejestracja do operacji

Atrybuty

Ops mogą mieć atrybuty, których wartości są ustawiane podczas dodawania operacji do wykresu. Są one używane do konfiguracji operacji, a ich wartości można uzyskać zarówno w implementacji jądra, jak iw typach wejść i wyjść w rejestracji operacji. Jeśli to możliwe, preferuj używanie danych wejściowych zamiast atrybutów, ponieważ dane wejściowe są bardziej elastyczne. Dzieje się tak, ponieważ attry są stałymi i muszą być zdefiniowane w czasie tworzenia wykresu. Natomiast dane wejściowe to tensory, których wartości mogą być dynamiczne; oznacza to, że dane wejściowe mogą zmieniać się na każdym kroku, być ustawiane za pomocą kanału informacyjnego itp. Atrybuty są używane do rzeczy, których nie można wykonać za pomocą danych wejściowych: dowolna konfiguracja, która wpływa na sygnaturę (liczba lub typ danych wejściowych lub wyjściowych) lub która może t zmieniać się krok po kroku.

Definiujesz attr podczas rejestrowania operacji, określając jego nazwę i typ przy użyciu metody Attr , która oczekuje specyfikacji 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 ZeroOut zachowała indeks określony przez użytkownika, zamiast tylko 0-tego elementu, możesz zarejestrować 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 danych wejściowych i wyjściowych).

Twoje jądro może wtedy uzyskać dostęp do tego attr 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 atrybucie obsługiwane są następujące typy:

  • string : Dowolna sekwencja bajtów (nie musi być UTF8).
  • int : liczba całkowita ze znakiem.
  • float : liczba zmiennoprzecinkowa.
  • bool : Prawda czy fałsz.
  • type : jedna z (nie-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 nieprawidłowy.

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

Domyślne wartości 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ącego <attr-type-expr> s:

{'<string1>', '<string2>'} : Wartość musi być ciągiem, który ma wartość <string1> lub <string2> . Nazwa typu, string , jest implikowana, gdy używasz tej składni. 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 implikowane, 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 : typ type ograniczony do typów numerycznych (nie łańcuchowych i nieliczbowych).
  • realnumbertype : Podobnie jak numbertype bez typów złożonych.
  • quantizedtype : Podobnie jak numbertype , ale tylko skwantowane typy liczb.

Konkretne listy dozwolonych przez nie typów są definiowane przez funkcje (takie jak NumberTypes() ) w tensorflow/core/framework/types.h . W tym przykładzie atrybut t musi być jednym z typów numerycznych:

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

Do 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 attr t być dowolnym typem liczbowym lub typem bool:

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

Do 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ć wartością int, której wartość jest większa lub równa <n> , gdzie <n> jest liczbą naturalną. Na przykład następująca rejestracja operacji określa, że 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 następująca rejestracja operacji określa, że a jest listą typów ( int32 lub float ) i że muszą być co najmniej 3 z nich:

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

Aby ustawić domyślną wartość atrybutu (co czyni go opcjonalnym w generowanym kodzie), dodaj = <default> na końcu, 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, jak określić wartość domyślną 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]");

Zwróć uwagę w szczególności, że wartości typu type używają tf.DType .

Wielopostaciowość

Typ polimorfizm

W przypadku operacji, które mogą przyjmować różne typy jako dane wejściowe lub wytwarzać różne typy danych wyjściowych, można określić atrybut w typie wejściowym lub wyjściowym w rejestracji operacji. Zazwyczaj należy wtedy zarejestrować OpKernel dla każdego obsługiwanego typu.

Na przykład, jeśli chcesz, aby ZeroOut działała na float s oprócz int32 s, twoja rejestracja operacji może wyglądać tak:

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

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

Nazywanie

Wejścia, wyjścia i atrybuty generalnie powinny mieć nazwy snake_case. Jedynym wyjątkiem są atrybuty, które są używane jako typ wejścia lub typ wyjścia. Te cechy można wywnioskować, gdy op zostanie dodany do wykresu, a więc nie pojawiają się w funkcji op. Na przykład ta ostatnia definicja ZeroOut wygeneruje funkcję Pythona, która wygląda tak:

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 , to T jest automatycznie ustawiane na int32 (właściwie DT_INT32 ). Te wywnioskowane atrybuty mają nazwy pisane wielkimi literami lub CamelCase.

Porównaj to z operacją, która ma atrybut typu określający typ wyjścia:

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 takim 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`.
  """
Przykład polimorfizmu typu
#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ą , podczas dodawania atrybutu do istniejącej operacji należy określić wartość domyślną :

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

Załóżmy, ż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 z nadmiarowym 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, operacje mogą zużywać lub wytwarzać zmienną liczbę tensorów.

W następnym przykładzie attr T zawiera listę typów i jest używany jako typ zarówno input in , jak i output out . Wejście i wyjście 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 dane wejściowe są listą tensorów float i double . Opera akceptuje na przykład typy wejściowe (float, double, float) iw takim przypadku typem wyjścia 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że to być również polimorficzne . W następnym przykładzie dane wejściowe to lista tensorów (o długości "N" ) tego samego (ale nieokreślonego) typu ( "T" ), a wyjściem 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ć tę wartość domyślną, używając ograniczenia ">=" w odpowiednim atr . W następnym przykładzie dane wejściowe to 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 operacji może mieć wiele wejść i wyjść:

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> to jedno 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> jest nazwą Attr z typem type lub list(type) (z możliwym ograniczeniem typu). Ta składnia pozwala na polimorficzne operacje .

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

    Odwołanie się do atrybutu typu list(type) umożliwia akceptację sekwencji 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");
    

    Zauważ, ż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 Attr o typie int . <type> może być albo tf.DType , albo nazwą atrybutu z typem type . Jako przykład pierwszego, ta operacja akceptuje listę tensorów int32 :

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

    Podczas gdy ta operacja 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")
    
  • Odwołanie do tensora: Ref(<type>) , gdzie <type> jest jednym z poprzednich typów.

Każdy atrybut użyty w typie danych wejściowych 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 chciałbyś w jakiś sposób wprowadzić zmiany w operacji.

Ogólnie rzecz biorąc, zmiany w istniejących, zaewidencjonowanych specyfikacjach muszą być zgodne z poprzednimi wersjami: zmiana specyfikacji operacji nie może zepsuć wcześniejszych serializowanych buforów protokołu GraphDef utworzonych na podstawie starszych specyfikacji. Szczegóły kompatybilności z GraphDefopisane tutaj .

Istnieje kilka sposobów zachowania zgodności z poprzednimi wersjami.

  1. Wszelkie nowe atrybuty dodawane do operacji muszą mieć zdefiniowane wartości domyślne, a przy tej wartości domyślnej operacja musi mieć oryginalne zachowanie. Aby zmienić operację z niepolimorficznej na polimorficzną, musisz nadać wartość domyślną nowemu typowi attr, 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 zgodny wstecz, używając:

    REGISTER_OP("MyGeneralUnaryOp")
        .Input("in: T")
        .Output("out: T")
        .Attr("T: numerictype = DT_FLOAT");
    
  2. Możesz bezpiecznie uczynić ograniczenie atrybutu mniej restrykcyjnym. 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 zmienić pojedyncze wejścia/wyjścia na wejścia/wyjścia listy, o ile domyślny typ listy jest zgodny ze starym podpisem.

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

  5. Przestrzeń nazw wszelkich nowych operacji, które tworzysz, poprzedzając nazwy operacji czymś unikalnym dla twojego projektu. Pozwala to uniknąć kolizji twojej operacji z jakimikolwiek operacjami, które mogą być zawarte w przyszłych wersjach TensorFlow.

  6. Planować naprzód! Spróbuj przewidzieć przyszłe zastosowania op. Niektórych zmian w podpisie nie można wykonać w sposób zgodny (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 dokonać zmiany w operacji, która jest zgodna z poprzednimi wersjami, utwórz nową operację o nowej nazwie z nową semantyką.

Należy również zauważyć, że chociaż te zmiany mogą zachować zgodność z GraphDef , wygenerowany kod Pythona może ulec zmianie w sposób, który nie jest zgodny ze starymi funkcjami wywołującymi. Interfejs API Pythona może być zachowany przez ostrożne zmiany w ręcznie napisanym opakowaniu Pythona, zachowując starą sygnaturę, z wyjątkiem ewentualnego dodawania nowych opcjonalnych argumentów na końcu. Generalnie niezgodne zmiany można wprowadzać tylko wtedy, gdy TensorFlow zmienia główne wersje i muszą być zgodne z semantyką wersji GraphDef .

Obsługa GPU

Możesz zaimplementować różne OpKernels i zarejestrować jeden dla procesora, a drugi dla GPU, tak 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/ . Zauważ, że niektóre jądra mają wersję procesora w pliku .cc , wersję GPU w pliku kończącym się na _gpu.cu.cc , a część kodu jest wspólna dla pliku .h .

Na przykład tf.pad zawiera 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 współdzielony kod to klasa szablonowa zdefiniowana w tensorflow/core/kernels/pad_op.h . Organizujemy kod w ten sposób z dwóch powodów: umożliwia to współdzielenie kodu przez implementacje procesora i GPU oraz umieszcza implementację GPU w oddzielnym pliku, dzięki czemu może być skompilowany tylko przez kompilator GPU.

Warto zauważyć, że nawet gdy używana jest wersja pad z jądrem GPU, nadal potrzebuje on swoich "paddings" w pamięci procesora. Aby zaznaczyć, że wejścia lub wyjścia są przechowywane w procesorze, dodaj 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 dla przykładu, który używa jądra CUDA do implementacji operacji. Biblioteka 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 ). Do użytku z binarną instalacją TensorFlow, jądra CUDA muszą być skompilowane za pomocą kompilatora NVIDIA nvcc . 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[@]}

Wytworzony powyżej cuda_op_kernel.so można załadować jak zwykle w Pythonie za pomocą funkcji tf.load_op_library .

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

Zaimplementuj gradient w Pythonie

Mając wykres operacji, TensorFlow używa automatycznego różnicowania (propagacji wstecznej), aby dodać nowe operacje reprezentujące gradienty w odniesieniu do istniejących operacji. Aby automatyczne różnicowanie działało dla nowych operacji, musisz zarejestrować funkcję gradientu, która oblicza gradienty w odniesieniu do wejść operacji, dane gradienty w odniesieniu do wyjść operacji.

Matematycznie, jeśli op obliczy \(y = f(x)\) , zarejestrowany gradient op przekształca gradienty \(\partial L/ \partial y\) straty \(L\) względem\(y\) na gradienty \(\partial L/ \partial x\) względem \(x\) za pomocą reguły łańcucha:

\[\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 wyjście, więc gradient względem wejścia jest rzadkim tensorem „jeden gorący”. 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 weźmie 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źć za pośrednictwem tf.Operation.get_attr .

  • Jeśli op ma wiele wyjść, funkcja gradientu przyjmie op i grads , gdzie grads jest listą 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 nie ma dobrze zdefiniowanego gradientu dla niektórych danych wejściowych, na przykład dla danych wejściowych liczb całkowitych używanych jako indeksy, odpowiedni zwracany gradient powinien mieć wartość None . Na przykład dla operacji przyjmującej tensor zmiennoprzecinkowy x i indeks całkowity i , funkcja gradientu return [x_grad, None] .

  • Jeśli nie ma żadnego znaczącego gradientu dla operacji, często nie będziesz musiał rejestrować żadnego gradientu i dopóki gradient operacji nigdy nie będzie potrzebny, wszystko będzie dobrze. W niektórych przypadkach operacja nie ma dobrze zdefiniowanego gradientu, ale może być zaangażowana w obliczanie gradientu. Tutaj możesz użyć ops.NotDifferentiable , aby automatycznie propagować zera wstecz.

Zauważ, że w momencie wywołania funkcji gradientowej dostępny jest tylko wykres przepływu danych operacji, a nie same dane tensora. W związku z tym wszystkie obliczenia muszą być wykonywane przy użyciu innych operacji tensorflow, aby mogły zostać uruchomione w czasie wykonywania grafu.

Funkcje kształtu w C++

Interfejs API TensorFlow ma funkcję o nazwie „wnioskowanie 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 REGISTER_OP języka C++ i pełnią dwie role: zapewnianie zgodności kształtów danych wejściowych podczas tworzenia wykresu oraz określanie kształtów danych wyjściowych.

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śli wyjście jest wybierane przez jego indeks, 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 danych wejściowych z indeksem idx można uzyskać przez c->input(idx) .

Istnieje wiele wspólnych funkcji kształtu, które mają zastosowanie do wielu operacji, takich jak shape_inference::UnchangedShape , które można znaleźć w common_shape_fns.h i 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. Dla wersji ZeroOut z ograniczeniem kształtu wektora 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 z dokładnie jednym wymiarem (lub jeśli kształt wejściowy jest nieznany, kształt wyjściowy będzie wektorem z jednym nieznanym wymiarem).

Jeśli twoja operacja jest polimorficzna z wieloma danymi wejściowymi , możesz użyć elementów członkowskich InferenceContext do określenia liczby kształtów do sprawdzenia i Merge do sprawdzenia, czy wszystkie kształty są zgodne (alternatywnie, uzyskaj dostęp do atrybutów, które wskazują długości, za pomocą InferenceContext::GetAttr , co 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ą się zmieniać dynamicznie, funkcje kształtu muszą być odporne na niekompletne informacje o kształcie dla dowolnych danych wejściowych. Metoda Merge w InferenceContext umożliwia obiektowi wywołującemu stwierdzenie, że dwa kształty są takie same, nawet jeśli jeden lub oba z nich 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 ma wiele funkcji, których można użyć do zdefiniowania manipulacji funkcjami kształtu. Na przykład możesz sprawdzić, czy określony 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 dla wszystkich różnych manipulacji 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. Możesz zobaczyć przykłady pisania tych testów w niektórych naszych testach core ops . (Składnia INFER_OK i INFER_ERROR jest trochę zagadkowa, ale staraj się być zwięzły w przedstawianiu specyfikacji kształtu wejściowego i wyjściowego w testach. Na razie zobacz otaczające komentarze w tych testach, aby poznać specyfikację ciągu kształtu).

Zbuduj pakiet pip dla swojej niestandardowej operacji

Aby zbudować pakiet pip dla swojej operacji, zobacz przykład tensorflow/custom-op . Ten przewodnik pokazuje, jak budować niestandardowe operacje z pakietu TensorFlow pip zamiast budować TensorFlow ze źródeł.