Tạo một op

Nếu bạn muốn tạo một op không có trong thư viện TensorFlow hiện có, chúng tôi khuyên bạn trước tiên nên thử viết op bằng Python dưới dạng thành phần của các op hoặc hàm Python hiện có. Nếu điều đó là không thể, bạn có thể tạo một C++ op tùy chỉnh. Có một số lý do khiến bạn có thể muốn tạo một C++ tùy chỉnh:

  • Thật không dễ dàng hoặc không thể thể hiện hoạt động của bạn như một sự kết hợp của các hoạt động hiện có.
  • Sẽ không hiệu quả nếu thể hiện hoạt động của bạn như một thành phần của các nguyên hàm hiện có.
  • Bạn muốn kết hợp thủ công một thành phần nguyên thủy mà trình biên dịch trong tương lai sẽ gặp khó khăn khi kết hợp.

Ví dụ: hãy tưởng tượng bạn muốn triển khai một cái gì đó như "tổng hợp trung vị", tương tự như toán tử "MaxPool", nhưng tính toán trung vị trên các cửa sổ trượt thay vì giá trị tối đa. Có thể thực hiện việc này bằng cách sử dụng tổ hợp các thao tác (ví dụ: sử dụng ExtractImagePatches và TopK), nhưng có thể không hiệu quả về hiệu suất hoặc bộ nhớ như thao tác gốc khi bạn có thể thực hiện điều gì đó thông minh hơn trong một thao tác hợp nhất duy nhất. Như thường lệ, trước tiên bạn nên cố gắng thể hiện điều mình muốn bằng cách sử dụng thành phần toán tử, chỉ chọn thêm một thao tác mới nếu việc đó tỏ ra khó khăn hoặc không hiệu quả.

Để kết hợp op tùy chỉnh của bạn, bạn sẽ cần phải:

  1. Đăng ký op mới trong tệp C++. Đăng ký op xác định một giao diện (đặc tả) cho chức năng của op, độc lập với việc triển khai của op. Ví dụ: đăng ký op xác định tên của op và đầu vào và đầu ra của op. Nó cũng định nghĩa hàm hình dạng được sử dụng để suy luận hình dạng tensor.
  2. Triển khai op trong C++. Việc triển khai op được gọi là hạt nhân và đó là cách triển khai cụ thể thông số kỹ thuật mà bạn đã đăng ký ở Bước 1. Có thể có nhiều hạt nhân cho các loại hoặc kiến ​​trúc đầu vào/đầu ra khác nhau (ví dụ: CPU, GPU).
  3. Tạo trình bao bọc Python (tùy chọn). Trình bao bọc này là API công khai được sử dụng để tạo op trong Python. Một trình bao bọc mặc định được tạo từ đăng ký op, có thể được sử dụng trực tiếp hoặc thêm vào.
  4. Viết hàm tính gradient cho op (tùy chọn).
  5. Kiểm tra hoạt động. Chúng tôi thường làm điều này bằng Python để thuận tiện, nhưng bạn cũng có thể kiểm tra op bằng C++. Nếu bạn xác định độ dốc, bạn có thể xác minh chúng bằng Python tf.test.compute_gradient_error . Xem relu_op_test.py làm ví dụ kiểm tra các hàm chuyển tiếp của các toán tử giống Relu và độ dốc của chúng.

Điều kiện tiên quyết

Xác định giao diện op

Bạn xác định giao diện của một op bằng cách đăng ký nó với hệ thống TensorFlow. Trong quá trình đăng ký, bạn chỉ định tên op của mình, đầu vào (loại và tên) và đầu ra (loại và tên), cũng như chuỗi tài liệu và bất kỳ attr nào mà op có thể yêu cầu.

Để xem cách thức hoạt động của nó, giả sử bạn muốn tạo một op có tensor là int32 s và xuất ra một bản sao của tensor, với tất cả trừ phần tử đầu tiên được đặt thành 0. Để thực hiện việc này, hãy tạo một tệp có tên zero_out.cc . Sau đó thêm lệnh gọi vào macro REGISTER_OP để xác định giao diện cho op của bạn:

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

Hoạt động ZeroOut này lấy một tensor to_zero của các số nguyên 32 bit làm đầu vào và xuất ra một tensor zeroed của các số nguyên 32 bit. Op cũng sử dụng hàm hình dạng để đảm bảo rằng tenxơ đầu ra có hình dạng giống với tenxơ đầu vào. Ví dụ: nếu đầu vào là một tensor có hình [10, 20] thì hàm hình dạng này chỉ định rằng hình dạng đầu ra cũng là [10, 20].

Triển khai kernel cho op

Sau khi bạn xác định giao diện, hãy cung cấp một hoặc nhiều cách triển khai op. Để tạo một trong những hạt nhân này, hãy tạo một lớp mở rộng OpKernel và ghi đè phương thức Compute . Phương thức Compute cung cấp một đối số context thuộc loại OpKernelContext* , từ đó bạn có thể truy cập những thứ hữu ích như các tensor đầu vào và đầu ra.

Thêm kernel của bạn vào tệp bạn đã tạo ở trên. Hạt nhân có thể trông giống như thế này:

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

Sau khi triển khai kernel, bạn đăng ký nó với hệ thống TensorFlow. Trong quá trình đăng ký, bạn chỉ định các ràng buộc khác nhau mà kernel này sẽ chạy. Ví dụ: bạn có thể có một hạt nhân được tạo cho CPU và một hạt nhân riêng cho GPU.

Để thực hiện việc này cho ZeroOut op, hãy thêm phần sau vào zero_out.cc :

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

Nhân CPU đa luồng

Để viết nhân CPU đa luồng, có thể sử dụng hàm Shard trong work_sharder.h . Hàm này phân chia chức năng tính toán trên các luồng được định cấu hình để sử dụng cho luồng nội bộ (xem inter_op_parallelism_threads trong config.proto ).

hạt nhân GPU

Nhân GPU được triển khai thành hai phần: hạt nhân OpKernel và hạt nhân CUDA cùng với mã khởi chạy của nó.

Đôi khi việc triển khai OpKernel là phổ biến giữa nhân CPU và GPU, chẳng hạn như kiểm tra đầu vào và phân bổ đầu ra. Trong trường hợp đó, cách triển khai được đề xuất là:

  1. Xác định khuôn mẫu OpKernel trên Thiết bị và loại nguyên thủy của tensor.
  2. Để thực hiện tính toán thực tế của đầu ra, hàm Điện toán gọi một cấu trúc functor theo khuôn mẫu.
  3. Chuyên môn hóa của functor đó cho CPUDevice được xác định trong cùng một tệp, nhưng chuyên môn hóa cho GPUDevice được xác định trong tệp .cu.cc, vì nó sẽ được biên dịch bằng trình biên dịch CUDA.

Đây là một ví dụ thực hiện. 

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

Xây dựng thư viện op

Biên dịch op bằng trình biên dịch hệ thống của bạn (cài đặt nhị phân TensorFlow)

Bạn có thể biên dịch zero_out.cc bằng trình biên dịch C++ như g++ hoặc clang có sẵn trên hệ thống của bạn. Gói PIP nhị phân cài đặt các tệp tiêu đề và thư viện mà bạn cần để biên dịch hoạt động của mình ở các vị trí dành riêng cho hệ thống. Tuy nhiên, thư viện python TensorFlow cung cấp hàm get_include để lấy thư mục tiêu đề và thư mục get_lib có một đối tượng dùng chung để liên kết. Đây là kết quả đầu ra của các chức năng này trên máy 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'

Giả sử bạn đã cài đặt g++ , đây là chuỗi lệnh bạn có thể sử dụng để biên dịch hoạt động của mình thành thư viện động.

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

Trên macOS, cần có cờ bổ sung "-không xác định Dynamic_lookup" khi xây dựng tệp .so .

Lưu ý về phiên bản gcc >=5 : gcc sử dụng C++ ABI mới kể từ phiên bản 5 . TensorFlow 2.8 trở về trước được xây dựng bằng gcc4 sử dụng ABI cũ hơn. Nếu bạn đang sử dụng các phiên bản TensorFlow này và đang cố gắng biên dịch thư viện op của mình bằng gcc>=5 , hãy thêm -D_GLIBCXX_USE_CXX11_ABI=0 vào dòng lệnh để làm cho thư viện tương thích với ABI cũ hơn. Theo mặc định, các gói TensorFlow 2.9+ tương thích với ABI mới hơn.

Biên dịch op bằng bazel (cài đặt nguồn TensorFlow)

Nếu bạn đã cài đặt nguồn TensorFlow, bạn có thể sử dụng hệ thống xây dựng của TensorFlow để biên dịch op. Đặt tệp BUILD tuân theo quy tắc xây dựng Bazel trong thư mục tensorflow/core/user_ops .

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

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

Chạy lệnh sau để xây dựng zero_out.so .

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

Để biên dịch thao tác Example , với CUDA Kernel, bạn cần sử dụng tham số gpu_srcs của tf_custom_op_library . Đặt tệp BUILD có quy tắc xây dựng Bazel sau vào một thư mục mới bên trong thư mục tensorflow/core/user_ops (ví dụ: "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"],
)

Chạy lệnh sau để xây dựng kernel_example.so .

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

Sử dụng op trong Python

API TensorFlow Python cung cấp hàm tf.load_op_library để tải thư viện động và đăng ký op với khung TensorFlow. load_op_library trả về một mô-đun Python chứa các trình bao bọc Python cho op và kernel. Do đó, khi đã xây dựng xong op, bạn có thể thực hiện các thao tác sau để chạy nó từ 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)

Hãy nhớ rằng hàm được tạo sẽ được đặt tên là snake_case (để tuân thủ PEP8 ). Vì vậy, nếu op của bạn có tên ZeroOut trong tệp C++, hàm python sẽ được gọi là zero_out .

Để làm cho op khả dụng dưới dạng một hàm thông thường có thể import từ mô-đun Python, có thể hữu ích khi thực hiện lệnh gọi load_op_library trong tệp nguồn Python như sau: 

import tensorflow as tf

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

Xác minh rằng op hoạt động

Một cách tốt để xác minh rằng bạn đã triển khai thành công hoạt động của mình là viết bài kiểm tra cho nó. Tạo tệp zero_out_op_test.py với nội dung:

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

Sau đó chạy thử nghiệm của bạn (giả sử bạn đã cài đặt tensorflow): 

$ python zero_out_op_test.py

Xây dựng các tính năng nâng cao vào hoạt động của bạn

Bây giờ bạn đã biết cách xây dựng một hoạt động và triển khai cơ bản (và có phần hạn chế), chúng ta sẽ xem xét một số điều phức tạp hơn mà bạn thường cần để xây dựng vào hoạt động của mình. Điều này bao gồm:

Kiểm tra và xác nhận có điều kiện

Ví dụ trên giả định rằng op áp dụng cho một tensor có hình dạng bất kỳ. Nếu nó chỉ áp dụng cho vectơ thì sao? Điều đó có nghĩa là thêm một bước kiểm tra vào quá trình triển khai OpKernel ở trên.

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

Điều này xác nhận rằng đầu vào là một vectơ và trả về việc đặt trạng thái InvalidArgument nếu không. Macro OP_REQUIRES có ba đối số:

Ngoài ra, nếu bạn muốn kiểm tra xem một đối tượng Status được trả về từ một số hàm có phải là lỗi hay không và nếu có thì hãy trả về nó, hãy sử dụng OP_REQUIRES_OK . Cả hai macro này đều trả về từ hàm bị lỗi.

Đăng ký hoạt động

Attr

Op có thể có attr, giá trị của chúng được đặt khi op được thêm vào biểu đồ. Chúng được sử dụng để định cấu hình op và các giá trị của chúng có thể được truy cập cả trong quá trình triển khai kernel cũng như trong các loại đầu vào và đầu ra trong đăng ký op. Ưu tiên sử dụng đầu vào thay vì attr khi có thể vì đầu vào linh hoạt hơn. Điều này là do attrs là hằng số và phải được xác định tại thời điểm xây dựng biểu đồ. Ngược lại, đầu vào là Tensor có giá trị động; nghĩa là, đầu vào có thể thay đổi từng bước, được đặt bằng nguồn cấp dữ liệu, v.v. Attr được sử dụng cho những việc không thể thực hiện được bằng đầu vào: bất kỳ cấu hình nào ảnh hưởng đến chữ ký (số lượng hoặc loại đầu vào hoặc đầu ra) hoặc có thể' t thay đổi theo từng bước.

Bạn xác định một attr khi đăng ký op, bằng cách chỉ định tên và loại của nó bằng phương thức Attr , phương thức này yêu cầu thông số kỹ thuật có dạng:

<name>: <attr-type-expr>

trong đó <name> bắt đầu bằng một chữ cái và có thể bao gồm các ký tự chữ và số và dấu gạch dưới, còn <attr-type-expr> là biểu thức kiểu có dạng được mô tả bên dưới .

Ví dụ: nếu bạn muốn ZeroOut op duy trì chỉ mục do người dùng chỉ định, thay vì chỉ phần tử thứ 0, bạn có thể đăng ký op như sau:

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

(Lưu ý rằng tập hợp các loại thuộc tính khác với tf.DType được sử dụng cho đầu vào và đầu ra.)

Sau đó, hạt nhân của bạn có thể truy cập attr này trong hàm tạo của nó thông qua tham số 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_;
};

sau đó có thể được sử dụng trong phương thức 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_);
  }

các loại attr

Các loại sau đây được hỗ trợ trong attr:

  • string : Bất kỳ chuỗi byte nào (không bắt buộc phải là UTF8).
  • int : Một số nguyên có dấu.
  • float : Một số dấu phẩy động.
  • bool : Đúng hay sai.
  • type : Một trong các giá trị (không phải ref) của DataType .
  • shape : Một TensorShapeProto .
  • list(<type>) : Danh sách <type> , trong đó <type> là một trong các loại trên. Lưu ý rằng list(list(<type>)) không hợp lệ.

Xem thêm: op_def_builder.cc:FinalizeAttr để biết danh sách chính xác.

Giá trị mặc định và ràng buộc

Attr có thể có giá trị mặc định và một số loại attr có thể có các ràng buộc. Để xác định một attr có ràng buộc, bạn có thể sử dụng các <attr-type-expr> sau:

{'<string1>', '<string2>'} : Giá trị phải là một chuỗi có giá trị <string1> hoặc <string2> . Tên của loại, string , được ngụ ý khi bạn sử dụng cú pháp này. Điều này mô phỏng một enum:

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

{<type1>, <type2>} : Giá trị thuộc loại type và phải là một trong <type1> hoặc <type2> , trong đó <type1> và <type2> được hỗ trợ tf.DType . Bạn không chỉ định rằng loại attr là type . Điều này được ngụ ý khi bạn có danh sách các loại trong {...} . Ví dụ: trong trường hợp này attr t là loại phải là int32 , float hoặc bool :

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

Có các phím tắt cho các ràng buộc loại phổ biến:

  • numbertype : type loại được giới hạn ở các loại số (không phải chuỗi và không phải bool).
  • realnumbertype : Giống như numbertype không có kiểu phức tạp.
  • quantizedtype : Giống như numbertype nhưng chỉ là loại số được lượng tử hóa.

Danh sách cụ thể các loại được các loại này cho phép được xác định bởi các hàm (như NumberTypes() ) trong tensorflow/core/framework/types.h . Trong ví dụ này attr t phải là một trong các kiểu số:

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

Đối với hoạt động này:

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

Danh sách có thể được kết hợp với các danh sách và loại đơn lẻ khác. Op sau cho phép attr t là bất kỳ loại số nào hoặc loại bool:

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

Đối với hoạt động này:

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> : Giá trị phải là int có giá trị lớn hơn hoặc bằng <n> , trong đó <n> là số tự nhiên. Ví dụ: đăng ký op sau chỉ định rằng attr a phải có giá trị ít nhất là 2 :

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

list(<type>) >= <n> : Danh sách loại <type> có độ dài lớn hơn hoặc bằng <n> . Ví dụ: đăng ký op sau chỉ định rằng attr a là danh sách các loại ( int32 hoặc float ) và phải có ít nhất 3 loại trong số đó:

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

Để đặt giá trị mặc định cho attr (làm cho nó tùy chọn trong mã được tạo), hãy thêm = <default> vào cuối, như trong:

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

Ngoài ra, cả ràng buộc và giá trị mặc định đều có thể được chỉ định:

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

Cú pháp được hỗ trợ của giá trị mặc định là cú pháp sẽ được sử dụng trong biểu diễn nguyên mẫu của định nghĩa GraphDef thu được.

Dưới đây là ví dụ về cách chỉ định mặc định cho tất cả các loại:

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

Đặc biệt lưu ý rằng các giá trị của loại type sử dụng tf.DType .

Đa hình

Kiểu đa hình

Đối với các hoạt động có thể lấy các loại khác nhau làm đầu vào hoặc tạo ra các loại đầu ra khác nhau, bạn có thể chỉ định attr trong loại đầu vào hoặc đầu ra trong đăng ký op. Thông thường, bạn sẽ đăng ký OpKernel cho từng loại được hỗ trợ.

Ví dụ: nếu bạn muốn ZeroOut op hoạt động trên float s ngoài int32 s, đăng ký op của bạn có thể trông như sau:

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

Đăng ký op của bạn hiện chỉ định rằng loại đầu vào phải là float hoặc int32 và đầu ra của nó sẽ có cùng loại vì cả hai đều có loại T .

Đặt tên

Thông thường, đầu vào, đầu ra và attr phải được đặt tên là snake_case. Một ngoại lệ là attr được sử dụng làm loại đầu vào hoặc loại đầu ra. Những attr đó có thể được suy ra khi op được thêm vào biểu đồ và do đó không xuất hiện trong hàm của op. Ví dụ: định nghĩa cuối cùng này của ZeroOut sẽ tạo ra một hàm Python trông như sau:

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

Nếu to_zero được truyền qua một tenxơ int32 thì T sẽ tự động được đặt thành int32 (thực tế là DT_INT32 ). Những attr được suy ra đó được đặt tên viết hoa hoặc CamelCase.

So sánh điều này với một op có loại attr xác định loại đầu ra:

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

Trong trường hợp này, người dùng phải chỉ định loại đầu ra, như trong Python được tạo: 

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`.
  """
Ví dụ về đa hình kiểu
#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);

Để duy trì khả năng tương thích ngược , bạn nên chỉ định giá trị mặc định khi thêm attr vào op hiện có:

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

Giả sử bạn muốn thêm nhiều loại hơn, giả sử double :

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

Thay vì viết một OpKernel khác với mã dự phòng như trên, bạn thường có thể sử dụng mẫu C++ để thay thế. Bạn vẫn sẽ có một đăng ký kernel ( lệnh gọi REGISTER_KERNEL_BUILDER ) cho mỗi lần quá tải.

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

Nếu bạn có nhiều hơn một vài lần quá tải, bạn có thể đặt đăng ký vào macro.

#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

Tùy thuộc vào danh sách loại bạn đang đăng ký kernel, bạn có thể sử dụng macro được cung cấp bởi 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
Liệt kê đầu vào và đầu ra

Ngoài việc có thể chấp nhận hoặc tạo ra các loại khác nhau, các op có thể tiêu thụ hoặc tạo ra một số lượng tensor khác nhau.

Trong ví dụ tiếp theo, attr T chứa danh sách các loại và được sử dụng làm loại của cả đầu in và đầu out . Đầu vào và đầu ra là danh sách các tensor thuộc loại đó (và số lượng cũng như loại tensor ở đầu ra giống như đầu vào, vì cả hai đều có loại T ).

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

Bạn cũng có thể đặt các hạn chế về loại có thể được chỉ định trong danh sách. Trong trường hợp tiếp theo này, đầu vào là danh sách các tensor float và double . Ví dụ, op chấp nhận các loại đầu vào (float, double, float) và trong trường hợp đó, loại đầu ra cũng sẽ là (float, double, float) .

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

Nếu bạn muốn tất cả các tensor trong danh sách có cùng loại, bạn có thể làm như sau:

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

Điều này chấp nhận danh sách các tensor int32 và sử dụng int attr N để chỉ định độ dài của danh sách.

Điều này cũng có thể được thực hiện theo kiểu đa hình . Trong ví dụ tiếp theo, đầu vào là danh sách các tensor (có độ dài "N" ) cùng loại (nhưng không xác định) ( "T" ) và đầu ra là một tensor đơn có loại khớp:

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

Theo mặc định, danh sách tensor có độ dài tối thiểu là 1. Bạn có thể thay đổi giá trị mặc định đó bằng cách sử dụng ràng buộc ">=" trên attr tương ứng . Trong ví dụ tiếp theo này, đầu vào là danh sách ít nhất 2 tensor int32 :

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

Cú pháp tương tự áp dụng cho attr "list(type)" : 

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

Đầu vào và đầu ra

Tóm lại những điều trên, đăng ký op có thể có nhiều đầu vào và đầu ra:

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

Mỗi thông số đầu vào hoặc đầu ra có dạng:

<name>: <io-type-expr>

trong đó <name> bắt đầu bằng một chữ cái và có thể bao gồm các ký tự chữ và số và dấu gạch dưới. <io-type-expr> là một trong các biểu thức kiểu sau:

  • <type> , trong đó <type> là loại đầu vào được hỗ trợ (ví dụ: float , int32 , string ). Điều này chỉ định một tensor đơn của loại đã cho.

    Xem tf.DType .

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

  • <attr-type> , trong đó <attr-type> là tên của Attr với loại type hoặc list(type) (có thể có hạn chế về loại). Cú pháp này cho phép các hoạt động đa hình .

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

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

    Tham chiếu attr của loại list(type) cho phép bạn chấp nhận một chuỗi các 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");

    Lưu ý rằng số lượng và loại tensor ở đầu out cũng giống như ở đầu vào in , vì cả hai đều thuộc loại T .

  • Đối với một chuỗi các tensor có cùng loại: <number> * <type> , trong đó <number> là tên của Attr có loại int . <type> có thể là tf.DType hoặc tên của attr có loại type . Như một ví dụ đầu tiên, op này chấp nhận danh sách các tensor int32 :

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

    Trong khi đó, op này chấp nhận danh sách các tensor thuộc bất kỳ loại nào, miễn là chúng đều giống nhau:

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

  • Để tham chiếu đến một tenxơ: Ref(<type>) , trong đó <type> là một trong các loại trước đó.

Bất kỳ attr nào được sử dụng trong loại đầu vào sẽ được suy ra. Theo quy ước, những attr được suy ra sử dụng tên viết hoa (như T hoặc N ). Mặt khác, đầu vào, đầu ra và attr có tên như tham số hàm (ví dụ: num_outputs ). Để biết thêm chi tiết, hãy xem phần trước về cách đặt tên .

Để biết thêm chi tiết, hãy xem tensorflow/core/framework/op_def_builder.h .

Khả năng tương thích ngược

Giả sử bạn đã viết một hoạt động tùy chỉnh, hay và chia sẻ nó với những người khác, nhờ đó bạn có được những khách hàng hài lòng khi sử dụng hoạt động của mình. Tuy nhiên, bạn muốn thực hiện các thay đổi đối với op theo một cách nào đó.

Nói chung, những thay đổi đối với các thông số kỹ thuật đã đăng ký hiện có phải tương thích ngược: việc thay đổi thông số kỹ thuật của một op không được phá vỡ bộ đệm giao thức GraphDef được tuần tự hóa trước đó được xây dựng từ các thông số kỹ thuật cũ hơn. Chi tiết về khả năng tương thích GraphDef được mô tả ở đây .

Có một số cách để duy trì khả năng tương thích ngược.

  1. Bất kỳ attr mới nào được thêm vào một thao tác đều phải có giá trị mặc định được xác định và với giá trị mặc định đó, op phải có hành vi ban đầu. Để thay đổi một thao tác từ không đa hình sang đa hình, bạn phải đặt giá trị mặc định cho kiểu attr mới để giữ nguyên chữ ký gốc theo mặc định. Ví dụ: nếu hoạt động của bạn là:

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

    bạn có thể làm cho nó đa hình theo cách tương thích ngược bằng cách sử dụng:

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

  2. Bạn có thể tạo một ràng buộc an toàn trên attr ít hạn chế hơn. Ví dụ: bạn có thể thay đổi từ {int32, int64} thành {int32, int64, float} hoặc type . Hoặc bạn có thể thay đổi từ {"apple", "orange"} thành {"apple", "banana", "orange"} hoặc string .

  3. Bạn có thể thay đổi đầu vào/đầu ra đơn thành đầu vào/đầu ra danh sách, miễn là mặc định cho loại danh sách khớp với chữ ký cũ.

  4. Bạn có thể thêm đầu vào/đầu ra danh sách mới nếu nó mặc định trống.

  5. Không gian tên cho bất kỳ hoạt động mới nào bạn tạo, bằng cách đặt trước tên hoạt động bằng nội dung nào đó duy nhất cho dự án của bạn. Điều này tránh việc hoạt động của bạn xung đột với bất kỳ hoạt động nào có thể được đưa vào các phiên bản tương lai của TensorFlow.

  6. Lên kế hoạch trước! Cố gắng dự đoán những ứng dụng trong tương lai của op. Một số thay đổi chữ ký không thể được thực hiện theo cách tương thích (ví dụ: tạo danh sách cùng loại thành danh sách có nhiều loại khác nhau).

Bạn có thể tìm thấy danh sách đầy đủ các thay đổi an toàn và không an toàn trong tensorflow/core/framework/op_compatibility_test.cc . Nếu bạn không thể thực hiện thay đổi đối với một thao tác tương thích ngược thì hãy tạo một thao tác mới với tên mới với ngữ nghĩa mới.

Cũng lưu ý rằng mặc dù những thay đổi này có thể duy trì khả năng tương thích với GraphDef nhưng mã Python được tạo có thể thay đổi theo cách không tương thích với các phương thức gọi cũ. API Python có thể được giữ tương thích bằng những thay đổi cẩn thận trong trình bao bọc Python viết tay, bằng cách giữ chữ ký cũ ngoại trừ việc có thể thêm các đối số tùy chọn mới vào cuối. Nói chung, những thay đổi không tương thích chỉ có thể được thực hiện khi TensorFlow thay đổi các phiên bản chính và phải tuân theo ngữ nghĩa của phiên bản GraphDef .

hỗ trợ GPU

Bạn có thể triển khai các OpKernels khác nhau và đăng ký một OpKernels cho CPU và một OpKernels khác cho GPU, giống như bạn có thể đăng ký kernel cho các loại khác nhau . Có một số ví dụ về hạt nhân có hỗ trợ GPU trong tensorflow/core/kernels/ . Lưu ý rằng một số hạt nhân có phiên bản CPU trong tệp .cc , phiên bản GPU trong tệp kết thúc bằng _gpu.cu.cc và một số mã được chia sẻ chung trong tệp .h .

Ví dụ: tf.pad có mọi thứ trừ nhân GPU trong tensorflow/core/kernels/pad_op.cc . Nhân GPU nằm trong tensorflow/core/kernels/pad_op_gpu.cu.cc và mã chia sẻ là một lớp được tạo khuôn mẫu được xác định trong tensorflow/core/kernels/pad_op.h . Chúng tôi sắp xếp mã theo cách này vì hai lý do: nó cho phép bạn chia sẻ mã chung giữa các triển khai CPU và GPU, đồng thời đặt việc triển khai GPU vào một tệp riêng để chỉ trình biên dịch GPU mới có thể biên dịch được.

Một điều cần lưu ý, ngay cả khi sử dụng phiên bản nhân GPU của pad , nó vẫn cần đầu vào "paddings" trong bộ nhớ CPU. Để đánh dấu rằng đầu vào hoặc đầu ra được giữ trên CPU, hãy thêm lệnh gọi HostMemory() vào đăng ký kernel, ví dụ: 

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

Biên dịch kernel cho thiết bị GPU

Hãy xem cuda_op_kernel.cu.cc để biết ví dụ sử dụng nhân CUDA để triển khai op. tf_custom_op_library chấp nhận đối số gpu_srcs trong đó có thể chỉ định danh sách các tệp nguồn chứa nhân CUDA (tệp *.cu.cc ). Để sử dụng với bản cài đặt nhị phân của TensorFlow, nhân CUDA phải được biên dịch bằng trình biên dịch nvcc của NVIDIA. Đây là chuỗi lệnh bạn có thể sử dụng để biên dịch cuda_op_kernel.cu.cc và cuda_op_kernel.cc thành một thư viện có thể tải động duy nhất:

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 được tạo ở trên có thể được tải như bình thường bằng Python, sử dụng hàm tf.load_op_library .

Lưu ý rằng nếu thư viện CUDA của bạn chưa được cài đặt trong /usr/local/lib64 , bạn sẽ cần chỉ định đường dẫn rõ ràng trong lệnh thứ hai (g++) ở trên. Ví dụ: thêm -L /usr/local/cuda-8.0/lib64/ nếu CUDA của bạn được cài đặt trong /usr/local/cuda-8.0 .

Triển khai độ dốc trong Python

Đưa ra một biểu đồ các op, TensorFlow sử dụng tính năng phân biệt tự động (lan truyền ngược) để thêm các op mới biểu thị độ dốc đối với các op hiện có. Để thực hiện phân biệt tự động cho các hoạt động mới, bạn phải đăng ký một hàm gradient tính toán độ dốc đối với đầu vào của các hoạt động đã cho độ dốc đối với đầu ra của hoạt động.

Về mặt toán học, nếu một op tính toán \(y = f(x)\) gradient op đã đăng ký chuyển đổi độ dốc \(\partial L/ \partial y\) mất mát \(L\) liên quan đến\(y\) thành độ dốc \(\partial L/ \partial x\) liên quan đến \(x\) thông qua quy tắc dây chuyền:

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

Trong trường hợp ZeroOut , chỉ có một mục trong đầu vào ảnh hưởng đến đầu ra, do đó độ dốc đối với đầu vào là một tenxơ "một nóng" thưa thớt. Điều này được thể hiện như sau:

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

Chi tiết về việc đăng ký hàm gradient với tf.RegisterGradient :

  • Đối với một op có một đầu ra, hàm gradient sẽ lấy một tf.Operation , op và một tf.Tensor grad và xây dựng các op mới từ các tenxơ op.inputs[i] , op.outputs[i] và grad . Thông tin về bất kỳ attr nào có thể được tìm thấy qua tf.Operation.get_attr .

  • Nếu op có nhiều đầu ra, hàm gradient sẽ lấy op và grads , trong đó grads là danh sách các gradient tương ứng với mỗi đầu ra. Kết quả của hàm gradient phải là danh sách các đối tượng Tensor biểu thị các gradient tương ứng với từng đầu vào.

  • Nếu không có độ dốc được xác định rõ ràng cho một số đầu vào, chẳng hạn như đối với đầu vào số nguyên được sử dụng làm chỉ mục, thì độ dốc trả về tương ứng sẽ là None . Ví dụ: đối với một op lấy tenxơ dấu phẩy động x và chỉ số nguyên i , hàm gradient sẽ return [x_grad, None] .

  • Nếu không có độ dốc có ý nghĩa nào cho op, bạn thường sẽ không phải đăng ký bất kỳ độ dốc nào và miễn là độ dốc của op không bao giờ cần thiết thì bạn sẽ ổn. Trong một số trường hợp, một op không có độ dốc được xác định rõ ràng nhưng có thể liên quan đến việc tính toán độ dốc. Tại đây bạn có thể sử dụng ops.NotDifferentiable để tự động truyền ngược các số 0.

Lưu ý rằng tại thời điểm hàm gradient được gọi, chỉ có biểu đồ luồng dữ liệu của ops chứ không có dữ liệu tensor. Do đó, tất cả tính toán phải được thực hiện bằng cách sử dụng các hoạt động tensorflow khác để chạy tại thời điểm thực thi biểu đồ.

Thêm gợi ý loại khi đăng ký gradient tùy chỉnh cho loại op để làm cho mã dễ đọc hơn, dễ sửa lỗi hơn, dễ bảo trì hơn và mạnh mẽ hơn thông qua xác thực dữ liệu. Ví dụ: khi lấy op làm tham số trong hàm, hãy chỉ định rằng hàm gradient sẽ lấy tf.Operation làm loại tham số.

Hàm hình dạng trong C++

API TensorFlow có một tính năng gọi là "suy luận hình dạng" cung cấp thông tin về hình dạng của tensor mà không cần phải thực thi biểu đồ. Suy luận hình dạng được hỗ trợ bởi "hàm hình dạng" được đăng ký cho từng loại op trong khai báo C++ REGISTER_OP và thực hiện hai vai trò: xác nhận rằng hình dạng của đầu vào tương thích trong quá trình xây dựng biểu đồ và chỉ định hình dạng cho đầu ra.

Các hàm hình dạng được định nghĩa là các thao tác trên lớp shape_inference::InferenceContext . Ví dụ: trong hàm hình dạng cho ZeroOut:

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

c->set_output(0, c->input(0)); tuyên bố rằng hình dạng của đầu ra đầu tiên phải được đặt thành hình dạng của đầu vào đầu tiên. Nếu đầu ra được chọn theo chỉ mục của nó như trong ví dụ trên, thì tham số thứ hai của set_output phải là đối tượng ShapeHandle . Bạn có thể tạo một đối tượng ShapeHandle trống bằng hàm tạo mặc định của nó. Đối tượng ShapeHandle cho đầu vào có chỉ mục idx có thể được lấy bằng c->input(idx) .

Có một số hàm hình dạng phổ biến áp dụng cho nhiều hoạt động, chẳng hạn như shape_inference::UnchangedShape có thể tìm thấy trong common_shape_fns.h và được sử dụng như sau:

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

Hàm hình dạng cũng có thể hạn chế hình dạng của đầu vào. Đối với phiên bản ZeroOut có ràng buộc hình dạng vector , hàm hình dạng sẽ như sau:

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

Lệnh gọi WithRank xác thực rằng hình dạng đầu vào c->input(0) có hình dạng có chính xác một chiều (hoặc nếu hình dạng đầu vào không xác định thì hình dạng đầu ra sẽ là một vectơ có một chiều không xác định).

Nếu op của bạn đa hình với nhiều đầu vào , bạn có thể sử dụng các thành viên của InferenceContext để xác định số lượng hình dạng cần kiểm tra và Merge để xác thực rằng tất cả các hình dạng đều tương thích (cách khác, truy cập các thuộc tính cho biết độ dài, với InferenceContext::GetAttr , cung cấp quyền truy cập vào các thuộc tính của 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();
    });

Vì suy luận hình dạng là một tính năng tùy chọn và hình dạng của tensor có thể thay đổi linh hoạt nên các hàm hình dạng phải mạnh mẽ để không có thông tin hình dạng không đầy đủ cho bất kỳ đầu vào nào. Phương thức Merge trong InferenceContext cho phép người gọi xác nhận rằng hai hình dạng giống nhau, ngay cả khi một hoặc cả hai hình dạng đó không có thông tin đầy đủ. Các hàm hình dạng được xác định cho tất cả các hoạt động cốt lõi của TensorFlow và cung cấp nhiều ví dụ sử dụng khác nhau.

Lớp InferenceContext có một số hàm có thể được sử dụng để xác định các thao tác hàm hình dạng. Ví dụ: bạn có thể xác thực rằng một thứ nguyên cụ thể có giá trị rất cụ thể bằng cách sử dụng InferenceContext::Dim và InferenceContext::WithValue ; bạn có thể chỉ định rằng thứ nguyên đầu ra là tổng/tích của hai thứ nguyên đầu vào bằng cách sử dụng InferenceContext::Add và InferenceContext::Multiply . Xem lớp InferenceContext để biết tất cả các thao tác hình dạng khác nhau mà bạn có thể chỉ định. Ví dụ sau đặt hình dạng của đầu ra đầu tiên thành (n, 3), trong đó đầu vào đầu tiên có hình dạng (n, ...)

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

Nếu bạn có một hàm hình dạng phức tạp, bạn nên cân nhắc việc thêm một thử nghiệm để xác nhận rằng các kết hợp hình dạng đầu vào khác nhau sẽ tạo ra các kết hợp hình dạng đầu ra mong đợi. Bạn có thể xem ví dụ về cách viết các bài kiểm thử này trong một số bài kiểm thử hoạt động cốt lõi của chúng tôi. (Cú pháp của INFER_OK và INFER_ERROR hơi khó hiểu, nhưng hãy cố gắng biểu diễn các thông số hình dạng đầu vào và đầu ra trong các thử nghiệm một cách nhỏ gọn. Hiện tại, hãy xem các nhận xét xung quanh trong các thử nghiệm đó để hiểu về đặc tả chuỗi hình dạng).

Xây dựng gói pip cho hoạt động tùy chỉnh của bạn

Để xây dựng gói pip cho op của bạn, hãy xem ví dụ tensorflow/custom-op . Hướng dẫn này chỉ ra cách xây dựng các hoạt động tùy chỉnh từ gói pip TensorFlow thay vì xây dựng TensorFlow từ nguồn.