Buat operasi

Jika Anda ingin membuat operasi yang tidak dicakup oleh pustaka TensorFlow yang sudah ada, sebaiknya Anda mencoba menulis operasi dengan Python terlebih dahulu sebagai komposisi operasi atau fungsi Python yang sudah ada. Jika itu tidak memungkinkan, Anda dapat membuat operasi C++ khusus. Ada beberapa alasan mengapa Anda mungkin ingin membuat operasi C++ khusus:

  • Tidak mudah atau tidak mungkin untuk menyatakan operasi Anda sebagai komposisi dari operasi yang ada.
  • Tidaklah efisien untuk menyatakan operasi Anda sebagai komposisi primitif yang ada.
  • Anda ingin memadukan komposisi primitif yang akan sulit digabungkan oleh kompiler masa depan.

Misalnya, bayangkan Anda ingin menerapkan sesuatu seperti "penggabungan median", mirip dengan operator "MaxPool", tetapi menghitung median melalui jendela geser, bukan nilai maksimum. Melakukan ini dengan menggunakan komposisi operasi dapat dilakukan (misalnya, menggunakan ExtractImagePatches dan TopK), tetapi mungkin tidak seefisien kinerja atau memori seperti operasi asli di mana Anda dapat melakukan sesuatu yang lebih pintar dalam satu operasi gabungan tunggal. Seperti biasa, biasanya pertama-tama ada baiknya mencoba mengungkapkan apa yang Anda inginkan menggunakan komposisi operator, hanya memilih untuk menambahkan operasi baru jika terbukti sulit atau tidak efisien.

Untuk menggabungkan op kustom Anda, Anda harus:

  1. Daftarkan op baru dalam file C++. Registrasi OP mendefinisikan antarmuka (spesifikasi) untuk fungsionalitas OP, yang tidak bergantung pada implementasi OP. Sebagai contoh, pendaftaran op mendefinisikan nama op dan input dan output op. Ini juga mendefinisikan fungsi bentuk yang digunakan untuk inferensi bentuk tensor.
  2. Terapkan op di C++. Implementasi op dikenal sebagai kernel, dan ini adalah implementasi konkret dari spesifikasi yang Anda daftarkan di Langkah 1. Mungkin ada beberapa kernel untuk jenis atau arsitektur input/output yang berbeda (misalnya, CPU, GPU).
  3. Buat pembungkus Python (opsional). Pembungkus ini adalah API publik yang digunakan untuk membuat op dengan Python. Pembungkus default dihasilkan dari pendaftaran op, yang dapat digunakan secara langsung atau ditambahkan ke.
  4. Tulis fungsi untuk menghitung gradien untuk op (opsional).
  5. Uji op. Kami biasanya melakukan ini dengan Python untuk kenyamanan, tetapi Anda juga dapat menguji op di C++. Jika Anda menentukan gradien, Anda dapat memverifikasinya dengan Python tf.test.compute_gradient_error . Lihat relu_op_test.py sebagai contoh yang menguji fungsi maju dari operator mirip Relu dan gradiennya.

Prasyarat

Tentukan antarmuka op

Anda menentukan antarmuka operasi dengan mendaftarkannya ke sistem TensorFlow. Dalam pendaftaran, Anda menentukan nama op Anda, inputnya (tipe dan nama) dan output (tipe dan nama), serta docstring dan attr apa pun yang mungkin diperlukan oleh op.

Untuk melihat cara kerjanya, misalkan Anda ingin membuat op yang mengambil tensor int32 s dan mengeluarkan salinan tensor, dengan semua kecuali elemen pertama disetel ke nol. Untuk melakukan ini, buat file bernama zero_out.cc . Kemudian tambahkan panggilan ke makro REGISTER_OP yang menentukan antarmuka untuk operasi Anda:

#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();
    });

Op ZeroOut ini mengambil satu tensor to_zero dari bilangan bulat 32-bit sebagai input, dan menghasilkan tensor zeroed dari bilangan bulat 32-bit. Op juga menggunakan fungsi bentuk untuk memastikan bahwa tensor keluaran memiliki bentuk yang sama dengan tensor masukan. Misalnya, jika input adalah tensor bentuk [10, 20], maka fungsi bentuk ini menentukan bentuk output juga [10, 20].

Terapkan kernel untuk op

Setelah Anda menentukan antarmuka, berikan satu atau lebih implementasi op. Untuk membuat salah satu kernel ini, buat kelas yang memperluas OpKernel dan mengganti metode Compute . Metode Compute menyediakan satu argumen context bertipe OpKernelContext* , yang darinya Anda dapat mengakses hal-hal berguna seperti tensor input dan output.

Tambahkan kernel Anda ke file yang Anda buat di atas. Kernel mungkin terlihat seperti ini:

#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);
  }
};

Setelah mengimplementasikan kernel Anda, Anda mendaftarkannya ke sistem TensorFlow. Dalam pendaftaran, Anda menentukan batasan yang berbeda di mana kernel ini akan dijalankan. Misalnya, Anda mungkin memiliki satu kernel yang dibuat untuk CPU, dan kernel terpisah untuk GPU.

Untuk melakukan ini untuk operasi ZeroOut , tambahkan berikut ini ke zero_out.cc :

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

Kernel CPU multi-utas

Untuk menulis kernel CPU multi-utas, fungsi Shard di work_sharder.h dapat digunakan. Fungsi ini memecah fungsi komputasi di seluruh utas yang dikonfigurasi untuk digunakan untuk threading intra-op (lihat intra_op_parallelism_threads di config.proto ).

kernel GPU

Kernel GPU diimplementasikan dalam dua bagian: OpKernel dan kernel CUDA serta kode peluncurannya.

Terkadang implementasi OpKernel adalah umum antara kernel CPU dan GPU, seperti seputar memeriksa input dan mengalokasikan output. Dalam hal ini, implementasi yang disarankan adalah untuk:

  1. Tentukan templat OpKernel pada Perangkat dan jenis tensor primitif.
  2. Untuk melakukan perhitungan output yang sebenarnya, fungsi Compute memanggil struct functor templated.
  3. Spesialisasi functor untuk CPUDevice tersebut ditentukan dalam file yang sama, tetapi spesialisasi untuk GPUDevice ditentukan dalam file .cu.cc, karena akan dikompilasi dengan kompiler CUDA.

Berikut adalah contoh penerapannya.

// 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

Membangun perpustakaan op

Kompilasi op menggunakan kompiler sistem Anda (instalasi biner TensorFlow)

Anda harus dapat mengkompilasi zero_out.cc dengan kompiler C++ seperti g++ atau clang yang tersedia di sistem Anda. Paket PIP biner menginstal file header dan pustaka yang Anda perlukan untuk mengkompilasi op Anda di lokasi yang spesifik untuk sistem. Namun, pustaka python TensorFlow menyediakan fungsi get_include untuk mendapatkan direktori header, dan direktori get_lib memiliki objek bersama untuk ditautkan. Berikut adalah output dari fungsi-fungsi ini pada mesin 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'

Dengan asumsi Anda telah menginstal g++ , berikut adalah urutan perintah yang dapat Anda gunakan untuk mengkompilasi op Anda menjadi pustaka dinamis.

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

Di macOS, flag tambahan "-undefined dynamic_lookup" diperlukan saat membuat file .so .

Catatan tentang versi gcc >=5 : gcc menggunakan ABI C++ baru sejak versi 5 . TensorFlow 2.8 dan sebelumnya dibuat dengan gcc4 yang menggunakan ABI lama. Jika Anda menggunakan versi TensorFlow ini dan mencoba mengompilasi pustaka op Anda dengan gcc>=5 , tambahkan -D_GLIBCXX_USE_CXX11_ABI=0 ke baris perintah agar pustaka kompatibel dengan ABI lama. Paket TensorFlow 2.9+ kompatibel dengan ABI yang lebih baru secara default.

Kompilasi op menggunakan bazel (instalasi sumber TensorFlow)

Jika Anda menginstal sumber TensorFlow, Anda dapat menggunakan sistem build TensorFlow untuk mengompilasi operasi Anda. Tempatkan file BUILD dengan mengikuti aturan build Bazel di direktori tensorflow/core/user_ops .

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

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

Jalankan perintah berikut untuk membangun zero_out.so .

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

Untuk mengompilasi operasi Example , dengan Kernel CUDA, Anda perlu menggunakan parameter gpu_srcs dari tf_custom_op_library . Tempatkan file BUILD dengan aturan build Bazel berikut di folder baru di dalam direktori tensorflow/core/user_ops (misalnya "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"],
)

Jalankan perintah berikut untuk membangun kernel_example.so .

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

Gunakan op dengan Python

TensorFlow Python API menyediakan fungsi tf.load_op_library untuk memuat library dinamis dan mendaftarkan op dengan framework TensorFlow. load_op_library mengembalikan modul Python yang berisi pembungkus Python untuk op dan kernel. Jadi, setelah Anda membuat op, Anda dapat melakukan hal berikut untuk menjalankannya dari Python:

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)

Perlu diingat, fungsi yang dihasilkan akan diberi nama snake_case (untuk mematuhi PEP8 ). Jadi, jika op Anda bernama ZeroOut di file C++, fungsi python akan dipanggil zero_out .

Untuk membuat op tersedia sebagai fungsi reguler import -able dari modul Python, mungkin berguna untuk memanggil load_op_library dalam file sumber Python sebagai berikut:

import tensorflow as tf

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

Verifikasi bahwa op bekerja

Cara yang baik untuk memverifikasi bahwa Anda telah berhasil mengimplementasikan operasi Anda adalah dengan menulis tes untuk itu. Buat file zero_out_op_test.py dengan isi:

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()

Kemudian jalankan pengujian Anda (dengan asumsi Anda telah menginstal tensorflow):

$ python zero_out_op_test.py

Bangun fitur lanjutan ke dalam operasi Anda

Sekarang setelah Anda mengetahui cara membangun op dan implementasi dasar (dan agak terbatas), kita akan melihat beberapa hal yang lebih rumit yang biasanya perlu Anda bangun ke dalam op Anda. Ini termasuk:

Pemeriksaan dan validasi bersyarat

Contoh di atas mengasumsikan bahwa op diterapkan pada tensor dalam bentuk apa pun. Bagaimana jika itu hanya diterapkan pada vektor? Itu berarti menambahkan centang pada implementasi OpKernel di atas.

  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."));
    // ...
  }

Ini menegaskan bahwa input adalah vektor, dan mengembalikan status InvalidArgument jika tidak. Makro OP_REQUIRES membutuhkan tiga argumen:

Alternatifnya, jika Anda ingin menguji apakah objek Status yang dikembalikan dari beberapa fungsi merupakan kesalahan, dan jika ya, kembalikan, gunakan OP_REQUIRES_OK . Kedua makro ini kembali dari fungsi pada kesalahan.

pendaftaran OP

Attr

Ops dapat memiliki attrs, yang nilainya ditetapkan saat op ditambahkan ke grafik. Ini digunakan untuk mengonfigurasi op, dan nilainya dapat diakses baik dalam implementasi kernel maupun dalam jenis input dan output dalam pendaftaran op. Lebih suka menggunakan input daripada attr jika memungkinkan, karena input lebih fleksibel. Ini karena attr adalah konstanta dan harus didefinisikan pada waktu pembuatan grafik. Sebaliknya, input adalah Tensor yang nilainya bisa dinamis; yaitu, input dapat mengubah setiap langkah, disetel menggunakan umpan, dll. Attr digunakan untuk hal-hal yang tidak dapat dilakukan dengan input: konfigurasi apa pun yang memengaruhi tanda tangan (jumlah atau jenis input atau output) atau yang dapat ' t berubah dari langkah ke langkah.

Anda mendefinisikan attr ketika Anda mendaftarkan op, dengan menentukan nama dan jenisnya menggunakan metode Attr , yang mengharapkan spek dari formulir:

<name>: <attr-type-expr>

di mana <name> dimulai dengan huruf dan dapat terdiri dari karakter alfanumerik dan garis bawah, dan <attr-type-expr> adalah ekspresi tipe dari bentuk yang dijelaskan di bawah ini .

Misalnya, jika Anda ingin op ZeroOut mempertahankan indeks yang ditentukan pengguna, alih-alih hanya elemen ke-0, Anda dapat mendaftarkan op seperti ini:

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

(Perhatikan bahwa set tipe atribut berbeda dari tf.DType yang digunakan untuk input dan output.)

Kernel Anda kemudian dapat mengakses attr ini di konstruktornya melalui parameter 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_;
};

yang kemudian dapat digunakan dalam metode 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_);
  }

jenis attr

Jenis berikut ini didukung dalam attr:

  • string : Setiap urutan byte (tidak harus UTF8).
  • int : Bilangan bulat bertanda.
  • float : Angka floating point.
  • bool : Benar atau salah.
  • type : Salah satu nilai (non-ref) dari DataType .
  • shape : A TensorShapeProto .
  • list(<type>) : Daftar <type> , di mana <type> adalah salah satu dari tipe di atas. Perhatikan bahwa list(list(<type>)) tidak valid.

Lihat juga: op_def_builder.cc:FinalizeAttr untuk daftar definitif.

Nilai dan batasan default

Attr mungkin memiliki nilai default, dan beberapa tipe attr mungkin memiliki batasan. Untuk mendefinisikan attr dengan batasan, Anda dapat menggunakan <attr-type-expr> s berikut:

{'<string1>', '<string2>'} : Nilai harus berupa string yang memiliki nilai <string1> atau <string2> . Nama tipe, string , tersirat saat Anda menggunakan sintaks ini. Ini mengemulasi enum:

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

{<type1>, <type2>} : Nilai bertipe type , dan harus salah satu dari <type1> atau <type2> , di mana <type1> dan <type2> didukung tf.DType . Anda tidak menentukan bahwa jenis attr adalah type . Ini tersirat ketika Anda memiliki daftar tipe di {...} . Misalnya, dalam hal ini attr t adalah tipe yang harus berupa int32 , float , atau bool :

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

Ada jalan pintas untuk batasan tipe umum:

  • numbertype : type tipe yang dibatasi untuk tipe numerik (non-string dan non-bool).
  • realnumbertype : Seperti numbertype tanpa tipe kompleks.
  • quantizedtype : Seperti numbertype tetapi hanya tipe angka terkuantisasi.

Daftar tipe spesifik yang diizinkan oleh ini ditentukan oleh fungsi (seperti NumberTypes() ) di tensorflow/core/framework/types.h . Dalam contoh ini attr t harus salah satu dari tipe numerik:

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

Untuk operasi ini:

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

Daftar dapat digabungkan dengan daftar lain dan tipe tunggal. Op berikut memungkinkan attr t menjadi salah satu dari tipe numerik, atau tipe bool:

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

Untuk operasi ini:

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> : Nilai harus berupa int yang nilainya lebih besar dari atau sama dengan <n> , di mana <n> adalah bilangan asli. Misalnya, pendaftaran op berikut menetapkan bahwa attr a harus memiliki nilai minimal 2 :

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

list(<type>) >= <n> : Daftar tipe <type> yang panjangnya lebih besar dari atau sama dengan <n> . Misalnya, pendaftaran op berikut menetapkan bahwa attr a adalah daftar tipe (baik int32 atau float ), dan setidaknya harus ada 3 tipe:

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

Untuk menetapkan nilai default untuk attr (menjadikannya opsional dalam kode yang dihasilkan), tambahkan = <default> di bagian akhir, seperti pada:

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

Selain itu, batasan dan nilai default dapat ditentukan:

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

Sintaks yang didukung dari nilai default adalah apa yang akan digunakan dalam representasi proto dari definisi GraphDef yang dihasilkan.

Berikut adalah contoh cara menentukan default untuk semua jenis:

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]");

Perhatikan secara khusus bahwa nilai dari type type menggunakan tf.DType .

Polimorfisme

Jenis polimorfisme

Untuk operasi yang dapat mengambil jenis yang berbeda sebagai masukan atau menghasilkan jenis keluaran yang berbeda, Anda dapat menentukan attr dalam jenis masukan atau keluaran dalam pendaftaran op. Biasanya Anda akan mendaftarkan OpKernel untuk setiap jenis yang didukung.

Misalnya, jika Anda ingin op ZeroOut bekerja pada float s selain int32 s, pendaftaran op Anda mungkin terlihat seperti:

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

Registrasi op Anda sekarang menetapkan bahwa tipe input harus float , atau int32 , dan outputnya akan bertipe sama, karena keduanya bertipe T .

Penamaan

Input, output, dan attr umumnya harus diberi nama snake_case. Satu-satunya pengecualian adalah attr yang digunakan sebagai tipe input atau tipe output. Attr tersebut dapat disimpulkan saat op ditambahkan ke grafik sehingga tidak muncul di fungsi op. Misalnya, definisi terakhir dari ZeroOut ini akan menghasilkan fungsi Python yang terlihat seperti:

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`.
  """

Jika to_zero dilewatkan tensor int32 , maka T secara otomatis diatur ke int32 (yah, sebenarnya DT_INT32 ). Attr yang disimpulkan tersebut diberi nama Kapitalisasi atau CamelCase.

Bandingkan ini dengan op yang memiliki tipe attr yang menentukan tipe keluaran:

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");

Dalam hal ini, pengguna harus menentukan tipe keluaran, seperti pada Python yang dihasilkan:

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`.
  """
Contoh polimorfisme tipe
#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);

Untuk mempertahankan kompatibilitas mundur , Anda harus menentukan nilai default saat menambahkan attr ke op yang ada:

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

Katakanlah Anda ingin menambahkan lebih banyak jenis, katakan double :

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

Alih-alih menulis OpKernel lain dengan kode redundan seperti di atas, sering kali Anda dapat menggunakan template C++. Anda masih memiliki satu registrasi kernel ( panggilan REGISTER_KERNEL_BUILDER ) per kelebihan beban.

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>);

Jika Anda memiliki lebih dari beberapa kelebihan, Anda dapat memasukkan pendaftaran ke dalam makro.

#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

Bergantung pada daftar jenis yang Anda daftarkan untuk kernel, Anda mungkin dapat menggunakan makro yang disediakan oleh 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
Daftar input dan output

Selain dapat menerima atau menghasilkan jenis yang berbeda, ops dapat mengkonsumsi atau menghasilkan sejumlah variabel tensor.

Pada contoh berikutnya, attr T menyimpan daftar tipe, dan digunakan sebagai tipe input in dan output out . Input dan output adalah daftar tensor dari tipe tersebut (dan jumlah serta tipe tensor di output sama dengan input, karena keduanya bertipe T ).

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

Anda juga dapat membatasi jenis apa yang dapat ditentukan dalam daftar. Dalam kasus berikutnya, inputnya adalah daftar tensor float dan double . op menerima, misalnya, tipe input (float, double, float) dan dalam hal ini tipe outputnya juga (float, double, float) .

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

Jika Anda ingin semua tensor dalam daftar bertipe sama, Anda dapat melakukan sesuatu seperti:

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

Ini menerima daftar tensor int32 , dan menggunakan int attr N untuk menentukan panjang daftar.

Ini bisa dibuat tipe polimorfik juga. Pada contoh berikutnya, inputnya adalah daftar tensor (dengan panjang "N" ) dari tipe yang sama (tetapi tidak ditentukan) ( "T" ), dan outputnya adalah tensor tunggal dengan tipe yang cocok:

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

Secara default, daftar tensor memiliki panjang minimal 1. Anda dapat mengubah default tersebut menggunakan batasan ">=" pada attr yang sesuai . Dalam contoh berikut ini, inputnya adalah daftar minimal 2 tensor int32 :

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

Sintaks yang sama bekerja dengan attrs "list(type)" :

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

Masukan dan keluaran

Untuk meringkas hal di atas, pendaftaran op dapat memiliki banyak input dan output:

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

Setiap spesifikasi input atau output berbentuk:

<name>: <io-type-expr>

di mana <name> dimulai dengan huruf dan dapat terdiri dari karakter alfanumerik dan garis bawah. <io-type-expr> adalah salah satu dari ekspresi tipe berikut:

  • <type> , di mana <type> adalah tipe input yang didukung (misalnya float , int32 , string ). Ini menentukan tensor tunggal dari tipe yang diberikan.

    Lihat tf.DType .

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

  • <attr-type> , di mana <attr-type> adalah nama Attr dengan type tipe atau list(type) (dengan kemungkinan batasan tipe). Sintaks ini memungkinkan untuk operasi polimorfik .

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

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

    Mereferensikan attr dari list(type) memungkinkan Anda untuk menerima urutan tensor.

    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");

    Perhatikan bahwa jumlah dan tipe tensor pada output out sama dengan input in , karena keduanya bertipe T .

  • Untuk urutan tensor dengan tipe yang sama: <number> * <type> , di mana <number> adalah nama Attr dengan tipe int . <type> bisa berupa tf.DType , atau nama attr dengan type type . Sebagai contoh yang pertama, op ini menerima daftar tensor int32 :

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

    Sedangkan op ini menerima daftar tensor jenis apa pun, asalkan semuanya sama:

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

  • Untuk referensi tensor: Ref(<type>) , di mana <type> adalah salah satu dari tipe sebelumnya.

Attr apa pun yang digunakan dalam jenis input akan disimpulkan. Dengan konvensi attr yang disimpulkan tersebut menggunakan nama kapital (seperti T atau N ). Jika tidak, masukan, keluaran, dan attr memiliki nama seperti parameter fungsi (mis. num_outputs ). Untuk detail lebih lanjut, lihat bagian sebelumnya tentang penamaan .

Untuk detail selengkapnya, lihat tensorflow/core/framework/op_def_builder.h .

Kompatibilitas mundur

Anggaplah Anda telah menulis operasi kustom yang bagus dan membaginya dengan orang lain, sehingga Anda memiliki pelanggan yang senang menggunakan operasi Anda. Namun, Anda ingin melakukan perubahan pada op dengan cara tertentu.

Secara umum, perubahan pada spesifikasi check-in yang sudah ada harus kompatibel dengan versi sebelumnya: mengubah spesifikasi op tidak boleh merusak buffer protokol GraphDef berseri sebelumnya yang dibangun dari spesifikasi yang lebih lama. Rincian kompatibilitas GraphDef dijelaskan di sini .

Ada beberapa cara untuk mempertahankan kompatibilitas mundur.

  1. Attr baru apa pun yang ditambahkan ke operasi harus memiliki nilai default yang ditentukan, dan dengan nilai default tersebut, op harus memiliki perilaku asli. Untuk mengubah operasi dari bukan polimorfik menjadi polimorfik, Anda harus memberikan nilai default ke attr tipe baru untuk mempertahankan tanda tangan asli secara default. Misalnya, jika operasi Anda adalah:

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

    Anda dapat membuatnya polimorfik dengan cara yang kompatibel dengan mundur menggunakan:

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

  2. Anda dapat dengan aman membuat batasan pada attr kurang ketat. Misalnya, Anda dapat mengubah dari {int32, int64} menjadi {int32, int64, float} atau type . Atau Anda dapat mengubah dari {"apple", "orange"} menjadi {"apple", "banana", "orange"} atau string .

  3. Anda dapat mengubah input/output tunggal menjadi input/output daftar, selama default untuk jenis daftar cocok dengan tanda tangan lama.

  4. Anda dapat menambahkan input / output daftar baru, jika defaultnya kosong.

  5. Namespace setiap operasi baru yang Anda buat, dengan mengawali nama operasi dengan sesuatu yang unik untuk proyek Anda. Ini menghindari op Anda bertabrakan dengan op apa pun yang mungkin disertakan dalam versi TensorFlow yang akan datang.

  6. Rencanakan ke depan! Cobalah mengantisipasi penggunaan op di masa mendatang. Beberapa perubahan tanda tangan tidak dapat dilakukan dengan cara yang kompatibel (misalnya, membuat daftar dengan jenis yang sama menjadi daftar dengan jenis yang berbeda-beda).

Daftar lengkap perubahan aman dan tidak aman dapat ditemukan di tensorflow/core/framework/op_compatibility_test.cc . Jika Anda tidak dapat mengubah operasi yang kompatibel mundur, buat operasi baru dengan nama baru dengan semantik baru.

Perhatikan juga bahwa meskipun perubahan ini dapat mempertahankan kompatibilitas GraphDef , kode Python yang dihasilkan dapat berubah dengan cara yang tidak kompatibel dengan pemanggil lama. API Python dapat tetap kompatibel dengan perubahan hati-hati dalam pembungkus Python yang ditulis tangan, dengan mempertahankan tanda tangan lama kecuali kemungkinan menambahkan argumen opsional baru ke bagian akhir. Perubahan yang umumnya tidak kompatibel hanya dapat dilakukan ketika TensorFlow mengubah versi utama, dan harus sesuai dengan semantik versi GraphDef .

Dukungan GPU

Anda dapat mengimplementasikan OpKernel yang berbeda dan mendaftarkan satu untuk CPU dan satu lagi untuk GPU, sama seperti Anda dapat mendaftarkan kernel untuk tipe yang berbeda . Ada beberapa contoh kernel dengan dukungan GPU di tensorflow/core/kernels/ . Perhatikan bahwa beberapa kernel memiliki versi CPU dalam file .cc , versi GPU dalam file yang diakhiri dengan _gpu.cu.cc , dan beberapa kode yang digunakan bersama dalam file .h .

Misalnya, tf.pad memiliki segalanya kecuali kernel GPU di tensorflow/core/kernels/pad_op.cc . Kernel GPU ada di tensorflow/core/kernels/pad_op_gpu.cu.cc , dan kode bersama adalah kelas template yang ditentukan di tensorflow/core/kernels/pad_op.h . Kami mengatur kode dengan cara ini karena dua alasan: memungkinkan Anda berbagi kode umum di antara implementasi CPU dan GPU, dan menempatkan implementasi GPU ke dalam file terpisah sehingga hanya dapat dikompilasi oleh kompiler GPU.

Satu hal yang perlu diperhatikan, bahkan ketika pad versi kernel GPU digunakan, masih membutuhkan input "paddings" di memori CPU. Untuk menandai bahwa input atau output disimpan di CPU, tambahkan panggilan HostMemory() ke registrasi kernel, misalnya:

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

Mengkompilasi kernel untuk perangkat GPU

Lihatlah cuda_op_kernel.cu.cc untuk contoh yang menggunakan kernel CUDA untuk mengimplementasikan op. tf_custom_op_library menerima argumen gpu_srcs di mana daftar file sumber yang berisi kernel CUDA ( file *.cu.cc ) dapat ditentukan. Untuk digunakan dengan instalasi biner TensorFlow, kernel CUDA harus dikompilasi dengan kompiler nvcc NVIDIA. Berikut adalah urutan perintah yang dapat Anda gunakan untuk mengkompilasi cuda_op_kernel.cu.cc dan cuda_op_kernel.cc menjadi satu pustaka yang dapat dimuat secara dinamis:

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 yang dihasilkan di atas dapat dimuat seperti biasa dengan Python, menggunakan fungsi tf.load_op_library .

Perhatikan bahwa jika pustaka CUDA Anda tidak dipasang di /usr/local/lib64 , Anda harus menentukan jalur secara eksplisit dalam perintah kedua (g++) di atas. Misalnya, tambahkan -L /usr/local/cuda-8.0/lib64/ jika CUDA Anda terpasang di /usr/local/cuda-8.0 .

Terapkan gradien dengan Python

Diberikan grafik operasi, TensorFlow menggunakan diferensiasi otomatis (perambatan balik) untuk menambahkan operasi baru yang mewakili gradien sehubungan dengan operasi yang ada. Agar diferensiasi otomatis berfungsi untuk operasi baru, Anda harus mendaftarkan fungsi gradien yang menghitung gradien sehubungan dengan input operasi yang diberikan gradien sehubungan dengan output operasi.

Secara matematis, jika op menghitung \(y = f(x)\) gradien terdaftar op mengonversi gradien \(\partial L/ \partial y\) dari kerugian \(L\) sehubungan dengan\(y\) menjadi gradien \(\partial L/ \partial x\) sehubungan dengan \(x\) melalui aturan rantai:

\[\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}.\]

Dalam kasus ZeroOut , hanya satu entri dalam input yang memengaruhi output, jadi gradien sehubungan dengan input adalah tensor "satu panas" yang jarang. Ini diungkapkan sebagai berikut:

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

Detail tentang mendaftarkan fungsi gradien dengan tf.RegisterGradient :

  • Untuk op dengan satu output, fungsi gradient akan mengambil tf.Operation , op , dan tf.Tensor grad dan membangun ops baru dari tensor op.inputs[i] , op.outputs[i] , dan grad . Informasi tentang attr apa pun dapat ditemukan melalui tf.Operation.get_attr .

  • Jika op memiliki banyak keluaran, fungsi gradien akan menggunakan op dan grads , di mana grads adalah daftar gradien sehubungan dengan setiap keluaran. Hasil dari fungsi gradien harus berupa daftar objek Tensor yang mewakili gradien sehubungan dengan setiap masukan.

  • Jika tidak ada gradien yang terdefinisi dengan baik untuk beberapa input, seperti untuk input bilangan bulat yang digunakan sebagai indeks, gradien yang dikembalikan harus None . Misalnya, untuk op yang menggunakan tensor titik mengambang x dan indeks bilangan bulat i , fungsi gradien akan return [x_grad, None] .

  • Jika tidak ada gradien yang berarti untuk op sama sekali, Anda seringkali tidak perlu mendaftarkan gradien apa pun, dan selama gradien op tidak diperlukan, Anda akan baik-baik saja. Dalam beberapa kasus, op tidak memiliki gradien yang terdefinisi dengan baik tetapi dapat dilibatkan dalam perhitungan gradien. Di sini Anda dapat menggunakan ops.NotDifferentiable untuk menyebarkan nol secara otomatis ke belakang.

Perhatikan bahwa saat fungsi gradien dipanggil, hanya grafik aliran data ops yang tersedia, bukan data tensor itu sendiri. Dengan demikian, semua perhitungan harus dilakukan menggunakan operasi tensorflow lainnya, untuk dijalankan pada waktu eksekusi grafik.

Tambahkan petunjuk tipe saat mendaftarkan gradien khusus untuk tipe op agar kode lebih mudah dibaca, dapat di-debug, lebih mudah dipelihara, dan lebih kuat melalui validasi data. Misalnya, saat menggunakan op sebagai parameter dalam suatu fungsi, tentukan bahwa fungsi gradien akan menggunakan tf.Operation sebagai tipe parameternya.

Bentuk fungsi dalam C++

TensorFlow API memiliki fitur yang disebut "inferensi bentuk" yang memberikan informasi tentang bentuk tensor tanpa harus mengeksekusi grafik. Inferensi bentuk didukung oleh "fungsi bentuk" yang didaftarkan untuk setiap jenis op dalam deklarasi C++ REGISTER_OP , dan melakukan dua peran: menyatakan bahwa bentuk input kompatibel selama pembuatan grafik, dan menentukan bentuk untuk output.

Fungsi bentuk didefinisikan sebagai operasi pada kelas shape_inference::InferenceContext . Misalnya, dalam fungsi bentuk untuk ZeroOut:

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

c->set_output(0, c->input(0)); menyatakan bahwa bentuk keluaran pertama harus diatur ke bentuk masukan pertama. Jika output dipilih berdasarkan indeksnya seperti pada contoh di atas, parameter kedua dari set_output harus berupa objek ShapeHandle . Anda dapat membuat objek ShapeHandle kosong dengan konstruktor defaultnya. Objek ShapeHandle untuk input dengan indeks idx dapat diperoleh dengan c->input(idx) .

Ada sejumlah fungsi bentuk umum yang berlaku untuk banyak operasi, seperti shape_inference::UnchangedShape yang dapat ditemukan di common_shape_fns.h dan digunakan sebagai berikut:

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

Fungsi bentuk juga dapat membatasi bentuk input. Untuk versi ZeroOut dengan batasan bentuk vektor , fungsi bentuknya adalah sebagai berikut:

    .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();
    });

Panggilan WithRank memvalidasi bahwa bentuk masukan c->input(0) memiliki bentuk dengan tepat satu dimensi (atau jika bentuk masukan tidak diketahui, bentuk keluaran akan berupa vektor dengan satu dimensi yang tidak diketahui).

Jika op Anda polimorfik dengan banyak masukan , Anda dapat menggunakan anggota InferenceContext untuk menentukan jumlah bentuk yang akan diperiksa, dan Merge untuk memvalidasi bahwa semua bentuk kompatibel (sebagai alternatif, akses atribut yang menunjukkan panjang, dengan InferenceContext::GetAttr , yang menyediakan akses ke atribut 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();
    });

Karena inferensi bentuk adalah fitur opsional, dan bentuk tensor dapat bervariasi secara dinamis, fungsi bentuk harus kuat untuk informasi bentuk yang tidak lengkap untuk input apa pun. Metode Merge di InferenceContext memungkinkan pemanggil untuk menyatakan bahwa dua bentuk adalah sama, meskipun salah satu atau keduanya tidak memiliki informasi yang lengkap. Fungsi bentuk ditentukan untuk semua operasi inti TensorFlow dan memberikan banyak contoh penggunaan yang berbeda.

Kelas InferenceContext memiliki sejumlah fungsi yang dapat digunakan untuk mendefinisikan manipulasi fungsi bentuk. Misalnya, Anda dapat memvalidasi bahwa dimensi tertentu memiliki nilai yang sangat spesifik menggunakan InferenceContext::Dim dan InferenceContext::WithValue ; Anda dapat menentukan bahwa dimensi output adalah jumlah / produk dari dua dimensi input menggunakan InferenceContext::Add dan InferenceContext::Multiply . Lihat kelas InferenceContext untuk berbagai manipulasi bentuk yang dapat Anda tentukan. Contoh berikut mengatur bentuk output pertama ke (n, 3), di mana input pertama memiliki bentuk (n, ...)

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

Jika Anda memiliki fungsi bentuk yang rumit, Anda harus mempertimbangkan untuk menambahkan pengujian untuk memvalidasi berbagai kombinasi bentuk masukan yang menghasilkan kombinasi bentuk keluaran yang diharapkan. Anda dapat melihat contoh cara menulis tes ini di beberapa tes operasi inti kami. (Sintaks INFER_OK dan INFER_ERROR sedikit samar, tetapi cobalah untuk kompak dalam merepresentasikan spesifikasi bentuk input dan output dalam pengujian. Untuk saat ini, lihat komentar di sekitarnya dalam pengujian tersebut untuk memahami spesifikasi string bentuk).

Bangun paket pip untuk operasi kustom Anda

Untuk membuat paket pip untuk operasi Anda, lihat contoh tensorflow/custom-op . Panduan ini menunjukkan cara membangun operasi khusus dari paket pip TensorFlow alih-alih membangun TensorFlow dari sumber.