Cảm ơn bạn đã theo dõi Google I/O. Xem tất cả các phiên theo yêu cầu Xem theo yêu cầu

Tạo một op

Nếu bạn muốn tạo một op không nằm 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 một 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 tùy chỉnh C++ op. Có một số lý do tại sao bạn có thể muốn tạo một tùy chỉnh C++ op:

  • Không dễ dàng hoặc không thể thể hiện hoạt động của bạn dưới dạng tổng hợp của các hoạt động hiện có.
  • Sẽ không hiệu quả khi thể hiện hoạt động của bạn dưới dạng 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à một 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 thứ 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ử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 một tập 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 năng hoặc bộ nhớ như một thao tác gốc, nơi 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, đơn lẻ. Như thường lệ, trước tiên bạn nên cố gắng thể hiện những gì bạn 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 điều đó 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 giao diện (thông số kỹ thuật) cho chức năng của op, độc lập với việc triển khai 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 một op được gọi là nhân và đó là việc triển khai cụ thể thông số kỹ thuật mà bạn đã đăng ký ở Bước 1. Có thể có nhiều 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. 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 toán độ dốc cho op (tùy chọn).
  5. Kiểm tra op. 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 trong 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 chức năng chuyển tiếp của các toán tử giống như 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 phần đăng ký, bạn chỉ định tên của op, đầu vào của nó (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 điều này hoạt động như thế nào, giả sử bạn muốn tạo một op nhận một tensor 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 đến 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();
    });

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

Triển khai hạt nhân cho op

Sau khi bạn xác định giao diện, hãy cung cấp một hoặc nhiều triển khai của 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ố ngữ context thuộc loại OpKernelContext* , từ đó bạn có thể truy cập những thứ hữu ích như tenxơ đầu vào và đầu ra.

Thêm hạt nhân 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 hạt nhân của mình, bạn đăng ký hạt nhân đó với hệ thống TensorFlow. Trong phần đăng ký, bạn chỉ định các ràng buộc khác nhau mà hạt nhân này sẽ chạy theo. Ví dụ: bạn có thể có một nhân được tạo cho CPU và một nhân riêng cho GPU.

Để thực hiện điều này cho op ZeroOut , 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

Để ghi nhân CPU đa luồng, có thể sử dụng hàm Shard trong work_sharder.h . Hàm này phân đoạn 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 intra_op_parallelism_threads trong config.proto ).

nhân GPU

Nhân GPU được triển khai thành hai phần: nhân OpKernel và nhân CUDA và 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ư xung quanh việc kiểm tra đầu vào và phân bổ đầu ra. Trong trường hợp đó, một triển khai được đề xuất là:

  1. Xác định khuôn mẫu OpKernel trên Thiết bị và kiểu 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 Tính toán gọi một cấu trúc hàm chức năng được tạo 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 sẽ có thể biên dịch zero_out.cc bằng trình biên dịch C++ chẳng hạn 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 tác vụ 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 tác vụ của mình thành một 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ờ bổ sung "-undefined dynamic_lookup" là bắt buộc khi tạo 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 với 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ố 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 các nguồn TensorFlow, bạn có thể sử dụng hệ thống xây dựng của TensorFlow để biên dịch tác vụ của mình. Đặt tệp BUILD với quy tắc xây dựng Bazel sau 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 Nhân CUDA, bạn cần sử dụng tham số gpu_srcs của tf_custom_op_library . Đặt tệp BUILD với 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. Vì vậy, khi bạn đã xây dựng 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)

Xin lưu ý 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 tên là ZeroOut trong các tệp C++, hàm python sẽ được gọi là zero_out .

Để làm cho op có sẵn dưới dạng hàm thông thường import - có thể nhập từ mô-đun Python, có thể hữu ích khi có 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 hay để xác minh rằng bạn đã triển khai thành công op của mình là viết mộ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 op 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 trong hoạt động của mình. Điêu nay 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 tenxơ 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 kiểm tra để 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 khẳng định rằng đầu vào là một vectơ và trả về đã đặt trạng thái InvalidArgument nếu không. Macro OP_REQUIRES nhận 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 hàm nào đó có phải là lỗi hay không và nếu có thì hãy trả lại nó, hãy sử dụng OP_REQUIRES_OK . Cả hai macro này đều trả về từ hàm do lỗi.

đăng ký op

Attrs

Các hoạt động có thể có attr, có giá trị được đặt khi hoạt động đượ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 hạt nhân và 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à Tensors có giá trị có thể động; nghĩa là, đầu vào có thể thay đổi từng bước, được đặt bằng cách sử dụng nguồn cấp dữ liệu, v.v. Attrs được sử dụng cho những thứ không thể thực hiện được với đầ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 từ bước này sang bước khác.

Bạn xác định một attr khi bạn đă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 mong đợi một 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, và <attr-type-expr> là một biểu thức loại có dạng được mô tả bên dưới .

Ví dụ: nếu bạn muốn op ZeroOut 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 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_;
};

mà 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 thu hút

Các loại sau được hỗ trợ trong một 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 đã ký.
  • float : Một số dấu phẩy động.
  • bool : Đúng hay sai.
  • type : Một trong những giá trị (không phải tham chiếu) của DataType .
  • shape : Một TensorShapeProto .
  • list(<type>) : Một 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.

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

Attrs có thể có các giá trị mặc định và một số loại attrs có thể có các ràng buộc. Để xác định một attr với các ràng buộc, bạn có thể sử dụng <attr-type-expr> s 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ó một danh sách các loại trong {...} . Ví dụ: trong trường hợp này, loại attr t 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 đượ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à các loại số được lượng tử hóa.

Danh sách cụ thể của các loại được phép bởi những điều này được xác định bởi các chức năng (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 op 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 khác và các loại đơn lẻ. Op sau đây 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 op 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ó trở thành 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ả giá trị 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à những gì sẽ được sử dụng trong biểu diễn nguyên mẫu của định nghĩa GraphDef kết quả.

Dưới đây là các 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 type loại sử dụng tf.DType .

đa hình

Loại đa hình

Đối với các op 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 một attr trong một loại đầu vào hoặc đầu ra trong đăng ký op. Thông thường, sau đó bạn sẽ đăng ký một OpKernel cho từng loại được hỗ trợ.

Chẳng hạn, nếu bạn muốn op ZeroOut hoạt động trên float ngoài int32 s, đăng ký op của bạn có thể giố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

Các đầu vào, đầu ra và attr thường phải được đặt tên Snake_case. Một ngoại lệ là attrs được sử dụng làm loại đầu vào hoặc loại đầu ra. Những attrs đó có thể được suy ra khi op được thêm vào biểu đồ và do đó không xuất hiện trong chức năng 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 có dạ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 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 luận đó đượ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 một 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, thường thì bạn sẽ có thể sử dụng mẫu C++ để thay thế. Bạn vẫn sẽ có một đăng ký hạt nhân ( REGISTER_KERNEL_BUILDER gọi) 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ý trong một 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 các 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 sản xuất các loại khác nhau, ops có thể tiêu thụ hoặc sản xuất một số lượng tenxơ khác nhau.

Trong ví dụ tiếp theo, attr T chứa một 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 tenxơ thuộc loại đó (và số lượng cũng như loại của các tenxơ ở đầu ra giống với đầ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 nào có thể được chỉ định trong danh sách. Trong trường hợp tiếp theo này, đầu vào là một danh sách các tenxơ 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 tenxơ trong danh sách cùng loại, bạn có thể làm điều gì đó như:

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

Điều này chấp nhận một danh sách các tenxơ 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 kiểu đa hình . Trong ví dụ tiếp theo, đầu vào là một danh sách các tenxơ (có độ dài "N" ) cùng loại (nhưng không xác định) ( "T" ) và đầu ra là một tenxơ đơn có loại phù hợ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 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 tenxơ int32 :

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

Cú pháp tương tự hoạt động với attrs "list(type)" :

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

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

Để tóm tắt những điều trên, một đă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ố kỹ thuật đầ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 tenxơ duy nhất 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 một Attr có type hoặc list(type) (có thể có hạn chế về loại). Cú pháp này cho phép 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 một attr của loại danh sách list(type) cho phép bạn chấp nhận một chuỗi các thang đo.

    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 tenxơ ở đầu out 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 tenxơ có cùng loại: <number> * <type> , trong đó <number> là tên của một Attr có loại int . <type> có thể là tf.DType hoặc tên của attr với type type . Như một ví dụ đầu tiên, op này chấp nhận một danh sách các tenxơ int32 :

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

    Trong khi đó, op này chấp nhận một danh sách các tenxơ thuộc bất kỳ loại nào, miễn là tất cả 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 tensor: 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 luận đó sử dụng tên viết hoa (như T hoặc N ). Mặt khác, đầu vào, đầu ra và attrs có tên giống như tham số chức năng (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 .

Tương thích ngược

Giả sử bạn đã viết một op tùy chỉnh, hay và chia sẻ nó với những người khác, vì vậy bạn 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 thay đổi op theo một cách nào đó.

Nói chung, các 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ỡ các bộ đệm giao thức GraphDef được đăng ký 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ả tại đây .

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

  1. Bất kỳ attrs mới nào được thêm vào một hoạt động 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 thành đa hình, bạn phải đặt giá trị mặc định cho loại mới attr để giữ nguyên chữ ký ban đầu theo mặc định. Ví dụ: nếu thao tác 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 trên attr ít hạn chế hơn một cách an toà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 phù hợp với chữ ký cũ.

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

  5. Không gian tên 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 một thứ gì đó duy nhất cho dự án của bạn. Điều này tránh để op của bạn va chạm với bất kỳ op 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 việc sử dụng op trong tương lai. 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 của mình đố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 xin lưu ý rằng mặc dù những thay đổi này có thể duy trì tính tương thích GraphDef , nhưng mã Python được tạo có thể thay đổi theo cách không tương thích với trình gọi cũ. API Python có thể được giữ tương thích bằng cách thay đổi cẩn thận trong trình bao bọc Python viết tay, bằng cách giữ nguyên chữ ký cũ ngoại trừ có thể thêm các đối số tùy chọn mới vào cuối. Nói chung, các 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 phù hợp với ngữ nghĩa của phiên bản GraphDef .

hỗ trợ GPU

Bạn có thể triển khai các hạt nhân OpKernel khác nhau và đăng ký một hạt nhân cho CPU và một hạt nhân khác cho GPU, giống như bạn có thể đăng ký hạt nhân 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ã dùng chung là một lớp khuôn mẫu được định nghĩa trong tensorflow/core/kernels/pad_op.h . Chúng tôi tổ chức 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 biệt để chỉ trình biên dịch GPU mới có thể biên dịch mã đó.

Một điều cần lưu ý, ngay cả khi phiên bản pad nhân GPU được sử dụng, nó vẫn cần đầu vào "paddings" của nó 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 một đối số gpu_srcs trong đó danh sách các tệp nguồn chứa nhân CUDA (tệp *.cu.cc ) có thể được chỉ định. Để sử dụng với bản cài đặt nhị phân của TensorFlow, các 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 trong Python, sử dụng hàm tf.load_op_library .

Lưu ý rằng nếu các thư viện CUDA của bạn chưa được cài đặt trong /usr/local/lib64 , thì bạn sẽ cần chỉ định rõ ràng đường dẫn 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 gradient trong Python

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

Về mặt toán học, nếu một op tính toán \(y = f(x)\) thì gradient op đã đăng ký sẽ chuyển đổi các gradient \(\partial L/ \partial y\) của \(L\) bị mất đối với\(y\) thành các gradient \(\partial L/ \partial x\) đối với \(x\) thông qua quy tắc chuỗi:

\[\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ỉ 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ý các hàm gradient với tf.RegisterGradient :

  • Đối với một op có một đầu ra, hàm gradient sẽ nhận 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, thì hàm gradient sẽ lấy op và grads , trong đó grads là danh sách các gradient đối với từng đầu ra. Kết quả của hàm gradient phải là một danh sách các đối tượng Tensor đại diện cho các gradient đối với từng đầu vào.

  • Nếu không có độ dốc được xác định rõ 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ỉ số, thì độ dốc được trả về tương ứng phải 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, bạn sẽ ổn thôi. Trong một số trường hợp, một op không có độ dốc được xác định rõ nhưng có thể tham gia vào việc tính toán độ dốc. Tại đây, bạn có thể sử dụng ops.NotDifferentiable để tự động chuyển các số không về phía sau.

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 phải 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, để được chạy tại thời điểm thực hiện biểu đồ.

Các hàm tạo hình 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 các tenxơ mà không cần phải thực thi biểu đồ. Suy luận hình dạng được hỗ trợ bởi "các 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ò: khẳng định 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 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, tham số thứ hai của set_output phải là một đố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ỉ số 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 thao tác, 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 vectơ , 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 WithRank xác thực rằng hình dạng đầu vào c->input(0) có hình dạng với chính xác một chiều (hoặc nếu hình dạng đầu vào không xác định, 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 cần kiểm tra và Merge để xác thực rằng tất cả các hình đề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 các tenxơ có thể thay đổi linh hoạt, nên các hàm hình dạng phải mạnh mẽ đối với 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 khẳng định rằng hai hình giống nhau, ngay cả khi một trong hai hoặc cả hai 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 của 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 một 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 xem xét thêm một thử nghiệm để xác thực 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 như mong đợi. Bạn có thể xem các ví dụ về cách viết các bài kiểm tra này trong một số bài kiểm tra 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 thể hiện các thông số kỹ thuật hình dạng đầu vào và đầu ra trong các bài kiểm tra một cách cô đọng. Hiện tại, hãy xem các nhận xét xung quanh trong các bài kiểm tra đó để hiểu về thông số kỹ thuật của chuỗi hình dạng).

Xây dựng gói pip cho op 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 op tùy chỉnh từ gói pip TensorFlow thay vì xây dựng TensorFlow từ nguồn.