Tạo một op

Sử dụng bộ sưu tập để sắp xếp ngăn nắp các trang Lưu và phân loại nội dung dựa trên lựa chọn ưu tiên của bạn.

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 hãy thử viết op bằng Python dưới dạng một thành phần của các chức năng hoặc hoạt động Python hiện có. Nếu không thể, bạn có thể tạo op C ++ tùy chỉnh. Có một số lý do tại sao bạn có thể muốn tạo một tùy chọn C ++ tùy chỉnh:

  • Không dễ hoặc không thể thể hiện hoạt động của bạn như một thành phần của các hoạt động hiện có.
  • Việc thể hiện hoạt động của bạn như một thành phần của các nguyên thủy hiện có là không hiệu quả.
  • Bạn muốn hợp nhất 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 hợp nhất.

Ví dụ: hãy tưởng tượng bạn muốn triển khai một cái gì đó như "gộp trung bình", tương tự như toán tử "MaxPool", nhưng tính toán trung bình qua cửa sổ trượt thay vì giá trị tối đa. Có thể thực hiện điều này bằng cách sử dụng một tổ hợp các hoạt động (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ư một hoạt động gốc, nơi bạn có thể thực hiện điều gì đó thông minh hơn trong một hoạt động được kết hợp duy nhất. Như mọi khi, 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ó hoặc không hiệu quả.

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

  1. Đăng ký op mới trong tệp C ++. Đăng ký op xác định một giao diện (đặc điểm kỹ thuật) cho chức năng của op, nó độ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à các đầu vào và đầu ra của op. Nó cũng xác định 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à nó là sự triển khai cụ thể của thông số kỹ thuật bạn đã đăng ký ở Bước 1. Có thể có nhiều hạt nhân cho các kiểu đầu vào / đầu ra hoặc kiến ​​trúc khác nhau (ví dụ: CPU, GPU).
  3. Tạo một 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 mộ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 lựa chọn trong C ++. Nếu bạn xác định gradient, bạn có thể xác minh chúng bằng Python tf.test.compute_gradient_error . Hãy 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 đă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ư docstrings và bất kỳ tập tin nào mà op có thể yêu cầu.

Để xem cách này hoạt động như thế nào, giả sử bạn muốn tạo một op lấy một 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 bằng 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 tùy chọn 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();
    });

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 zeroed bằng 0 của số nguyên 32 bit. Op cũng sử dụng một chức năng hình dạng để đả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 tensor có hình dạng [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 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 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 kiểu OpKernelContext* , từ đó bạn có thể truy cập những thứ hữu ích như các 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. Kernel có thể trông giống như sau:

#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 thực hiện kernel, bạn đăng ký nó với hệ thống TensorFlow. Trong đă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. Ví dụ: bạn có thể có một nhân được tạo cho CPU và một nhân riêng biệt dành cho GPU.

Để thực hiện việc 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

Để viết mộ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 đoạn một hàm tính toán trên các luồng được định cấu hình để sử dụng cho phân luồng nội bộ (xem intra_op_parallelism_threads trong config.proto ).

Hạt nhân GPU

Nhân GPU được thực hiện thành hai phần: Nhân OpKernel và CUDA và mã khởi chạy của nó.

Đôi khi việc triển khai OpKernel diễn ra 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 OpKernel templated trên Device 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 Compute gọi một cấu trúc hàm tạo mẫu.
  3. Chuyên môn hóa của bộ chức năng đó cho CPUDevice được định nghĩa 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ùy chọn của mình tại 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 được chia sẻ để liên kết chống lại. Đâ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 op 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 . Các gói pip nhị phân có sẵn trên trang web TensorFlow được xây dựng bằng gcc4 sử dụng ABI cũ hơn. Nếu bạn biên dịch thư viện op của mình với 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.

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 op của mình. Đặt tệp BUILD với quy tắc xây dựng Bazel sau đây 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 hoạt động 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 đây trong 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

TensorFlow Python API cung cấp hàm tf.load_op_library để tải thư viện động và đăng ký tùy chọn với khung TensorFlow. load_op_library trả về một mô-đun Python có chứa các trình bao bọc Python cho op và kernel. Do đó, khi bạn đã xây dựng op, bạn có thể làm như 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ên là solid_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 ++, thì 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ột mô-đun Python, có thể hữu ích nếu 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 tốt để xác minh rằng bạn đã triển khai thành công lựa chọn 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 hoạt động của bạn

Bây giờ bạn đã biết cách xây dựng và triển khai op cơ bản (và hơi 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 op của mình. Điêu nay bao gôm:

Kiểm tra có điều kiện và xác nhậ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ỳ. Điều gì sẽ xảy ra nếu nó chỉ áp dụng cho vectơ? Điều đó có nghĩa là thêm một dấu kiểm vào việc 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ề sau khi đặt InvalidArgument thái Đối số không hợp lệ nếu nó không phải là. Macro OP_REQUIRES có ba đối số:

Ngoài ra, nếu bạn muốn kiểm tra xem đố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ó trả về nó, hãy sử dụng OP_REQUIRES_OK . Cả hai macro này đều trả về từ hàm khi bị lỗi.

Đăng ký op

Đính kèm

Ops có thể có các phần đính kèm, có giá trị được đặt khi op được thêm vào biểu đồ. Chúng được sử dụng để 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à 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ì tập tin đính kèm khi có thể, vì đầu vào linh hoạt hơn. Điều này là do tập tin là hằng số và phải được xác định tại thời điểm xây dựng đồ thị. Ngược lại, đầu vào là các Tensors mà giá trị của nó có thể là động; nghĩa là, đầu vào có thể thay đổi mọi bước, được đặt bằng nguồn cấp dữ liệu, v.v. Phần đính kèm được sử dụng cho những việc không thể thực hiện 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à kiểu của nó bằng cách sử dụng phương thức Attr , phương thức này yêu cầu một đặc tả của biểu mẫu:

<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 kiểu của biểu mẫu được mô tả bên dưới .

Ví dụ: nếu bạn muốn 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ý tùy chọn 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 các đầu vào và đầu ra.)

Sau đó, hạt nhân của bạn có thể truy cập tệp đính kèm này trong phương thức khởi 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_);
  }

Attr các loại

Các loại sau được hỗ trợ trong một tập tin đính kèm:

  • 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 các <type> , trong đó <type> là một trong các kiểu 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 cuối cùng.

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

Phần đính kèm có thể có giá trị mặc định và một số loại phần đính kèm có thể có các ràng buộc. Để xác định tệp đính kèm 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 kiểu, 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><type2> được hỗ trợ tf.DType . Bạn không chỉ định rằng loại tệp đính kèm 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à một kiểu 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 loại ràng buộc phổ biến:

  • numbertype : Loại type bị hạn chế đối với 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.
  • kiểu quantizedtype : Giống numbertype nhưng chỉ là kiểu số lượng tử hóa.

Danh sách cụ thể của các loại được cho phép bởi các hàm này đượ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 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 cho phép attr t là bất kỳ kiểu số nào hoặc kiểu 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 kiểu <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 trong số chúng:

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

Để đặt giá trị mặc định cho một phần tử đính kèm (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 proto 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 loại type sử dụng tf.DType .

Tính đa hình

Nhập đ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 một tệp tin đính kèm 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 mỗi loại được hỗ trợ.

Ví dụ: nếu bạn muốn op ZeroOut hoạt động trên float s ngoài int32 s, đăng ký op của bạn có thể trông 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 bây giờ chỉ định rằng kiểu của đầu vào phải là float , hoặc int32 và đầu ra của nó sẽ là cùng một kiểu, vì cả hai đều có kiểu T

Đặt tên

Đầu vào, đầu ra và tập tin đính kèm thường phải được đặt tên solid_case. Một ngoại lệ là các tập tin được sử dụng làm kiểu đầu vào hoặc kiểu đầu ra. Những chứng nhận đó 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 op. Ví dụ: định nghĩa cuối cùng này về ZeroOut sẽ tạo ra một hàm Python giố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 tensor int32 , thì T sẽ tự động được đặt thành int32 (tốt, thực tế là DT_INT32 ). Các chứng từ đượ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ó kiểu attr xác định kiểu đầ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 kiểu đầ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`.
  """
Nhập ví dụ về đa hình
#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 tệp đính kèm vào tùy chọn 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 double hơn, giả sử như sau:

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ột mẫu C ++ để thay thế. Bạn sẽ vẫn có một đăng ký hạt nhân (cuộc 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ý 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ý hạt nhân, 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, các hoạt động có thể tiêu thụ hoặc sản xuất một số lượng tenxơ thay đổi.

Trong ví dụ tiếp theo, attr T chứa một danh sách các kiểu và được sử dụng làm kiểu 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 và loại tenxơ trong đầu ra giống với đầu vào, vì cả hai đều có kiểu T ).

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

Bạn cũng có thể đặt giới hạn về những 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 floatdouble tensors. Ví dụ, op chấp nhận các kiểu đầu vào (float, double, float) và trong trường hợp đó, kiểu đầ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 thuộc cùng một 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 một danh sách các tensors int32 và sử dụng một 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 loại đa hình . Trong ví dụ tiếp theo, đầu vào là một danh sách các tenxơ (với độ dài "N" ) của cùng một loại (nhưng không xác định) ( "T" ) và đầu ra là một tenxơ đơn của kiểu so 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 có í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 "list(type)" attrs:

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

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

Để tóm tắt ở 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ố đầ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 những biểu thức kiểu sau:

  • <type> , trong đó <type> là kiểu đầu vào được hỗ trợ (ví dụ: float , int32 , string ). Điều này chỉ định một tensor 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 Attr có kiểu type hoặc list(type) (có thể có hạn chế về kiểu). 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 đến một list(type) attr cho phép bạn chấp nhận một chuỗi các tenxơ.

    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 với đầ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 kiểu: <number> * <type> , trong đó <number> là tên của một Attr có kiểu int . <type> có thể là tf.DType hoặc tên của tệp đính kèm với kiểu type . Như một ví dụ về đầu tiên, op này chấp nhận một danh sách các tensors 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")
    
  • Đối với tham chiếu đến tensor: Ref(<type>) , trong đó <type> là một trong các kiểu trước đó.

Bất kỳ tập tin đính kèm nào được sử dụng trong loại đầu vào sẽ được suy ra. Theo quy ước, những phần tử được suy ra này sử dụng tên viết hoa (như T hoặc N ). Nếu không thì các đầu vào, đầu ra và tập tin đính kèm có tên giống 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 .

Tương thích ngược

Giả sử bạn đã viết một op tùy chỉnh, đẹp mắt 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 thực hiện các thay đổi đối với tùy chọn theo một cách nào đó.

Nói chung, các thay đổi đối với 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 của 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. Mọi tập tin đính kèm mới được thêm vào một thao tác 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 cung cấp giá trị mặc định cho kiểu đính kèm mới để giữ nguyên chữ ký ban đầu 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ể thực hiện một cách an toàn hạn chế đối với một tập tin í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 lẻ thành đầu vào / đầu ra danh sách, miễn là giá trị mặc định cho loại danh sách khớ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 danh sách đó được đặt mặc định là 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 thêm tiền tố tên hoạt động với một cái 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ỳ hoạt động nào có thể có trong các phiên bản TensorFlow trong tương lai.

  6. Lên kế hoạch trước! Cố gắng dự đoán các mục đích sử dụng trong tương lai cho op. Không thể thực hiện một số thay đổi chữ ký theo cách tương thích (ví dụ: tạo danh sách cùng loại thành danh sách các 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 hoạt động tương thích ngược, thì hãy tạo một hoạt động 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 trình gọi cũ. API Python có thể được giữ tương thích bằng các 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, 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 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 OpKernel khác nhau và đăng ký một nhân cho CPU và một nhân khác cho GPU, giống như bạn có thể đăng ký các 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ã chia sẻ là một lớp 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 quá trình triển khai CPU và GPU, đồng thời đặt quá trình triển khai GPU vào một tệp riêng biệt để nó chỉ có thể được biên dịch bởi trình biên dịch GPU.

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" trong bộ nhớ CPU. Để đánh dấu rằng các đầ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ý hạt nhân, 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 hạt nhân 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. Đối tf_custom_op_library chấp nhận đối số gpu_srcs trong đó danh sách tệp nguồn chứa nhân CUDA (tệp *.cu.cc ) có thể được chỉ định. Để sử dụng với 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. Dưới đây là chuỗi lệnh bạn có thể sử dụng để biên dịch cuda_op_kernel.cu.cccuda_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, bằng cách sử dụng hàm tf.load_op_library .

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

Đưa ra một biểu đồ gồm các hoạt động, TensorFlow sử dụng phân biệt tự động (backpropagation) để thêm các hoạt động mới đại diện cho độ dốc so với các hoạt động hiện có. Để làm cho sự phân biệt tự động hoạ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 các độ dốc liên quan đến các đầu vào của hoạt động đã cho các gradient liên quan đến đầu ra của các hoạt động.

Về mặt toán học, nếu op tính toán \(y = f(x)\) thì gradient op đã đăng ký sẽ chuyển đổi gradient \(\partial L/ \partial y\) mất \(L\) đối với\(y\) thành 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 của ZeroOut , chỉ một mục trong đầu vào ảnh hưởng đến đầu ra, do đó, gradient đối với đầu vào là tensor "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

Thông tin chi tiết về việc đăng ký các hàm gradient với tf.RegisterGradient :

  • Đối với op có một đầu ra, hàm gradient sẽ nhận tf.Operation , op , và tf.Tensor grad và xây dựng các hoạt động mới từ tensors op.inputs[i] , op.outputs[i]grad . Thông tin về bất kỳ tập tin đính kèm 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ẽ nhận opgrads , trong đó grads là danh sách các gradient liên quan đến mỗi đầu ra. Kết quả của hàm gradient phải là danh sách các đối tượng Tensor đại diện cho các gradient đối với mỗi đầu vào.

  • Nếu không có gradient được xác định rõ ràng cho một số đầu vào, chẳng hạn như cho các đầu vào số nguyên được sử dụng làm chỉ số, thì gradient trả về tương ứng phải là None . Ví dụ: đối với op lấy 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ó gradient nào có ý nghĩa cho op, bạn thường sẽ không phải đăng ký bất kỳ gradient nào và miễn là không bao giờ cần đến gradient của op, bạn sẽ ổn. Trong một số trường hợp, một op không có gradient được xác định rõ ràng nhưng có thể tham gia vào việc tính toán gradient. Ở đây bạn có thể sử dụng ops.NotDifferentiable để tự động truyền ngược các số không.

Lưu ý rằng tại thời điểm gọi hàm gradient, chỉ có biểu đồ luồng dữ liệu của ops, không có sẵn dữ liệu tensor. Do đó, tất cả cá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 gian thực hiện đồ thị.

Các hàm định dạng trong C ++

API TensorFlow có một tính năng được 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 hiện 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 đồ thị 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 hoạt động 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 bởi chỉ mục của nó như trong ví dụ trên, thì 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 phương thức khởi 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ể được 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);

Một 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 gọi WithRank xác nhận 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 là 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 là đ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 các dây căng có thể thay đổi linh hoạt, các chức năng 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 dạng giống nhau, ngay cả khi một trong hai hoặc cả hai không có thông tin đầy đủ. Các chức năng hình dạng được xác định cho tất cả các hoạt động TensorFlow cốt lõi 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 với hàm hình dạng. Ví dụ: bạn có thể xác nhận rằng một thứ nguyên cụ thể có giá trị rất cụ thể bằng cách sử dụng InferenceContext::DimInferenceContext::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::AddInferenceContext::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 chức năng hình dạng phức tạp, bạn nên xem xét thêm một bài kiểm tra để xác nhận rằng các kết hợp hình dạng đầu vào khác nhau tạo ra các kết hợp hình dạng đầu ra 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_OKINFER_ERROR hơi khó hiểu, nhưng hãy cố gắng gọn nhẹ trong việc biểu diễn các thông số kỹ thuật hình dạng đầu vào và đầu ra trong các thử nghiệm. 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 điểm chuỗi hình dạng).

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

Để xây dựng một gói pip cho op của bạn, hãy xem ví dụ về 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.