增加一个新 Op

预备知识:

如果现有的库没有涵盖你想要的操作, 你可以自己定制一个. 为了使定制的 Op 能够兼容原有的库
, 你必须做以下工作:

  • 在一个 C++ 文件中注册新 Op. Op 的注册与实现是相互独立的. 在其注册时描述了 Op 该如何执行.
    例如, 注册 Op 时定义了 Op 的名字, 并指定了它的输入和输出.
  • 使用 C++ 实现 Op. 每一个实现称之为一个 “kernel”, 可以存在多个 kernel, 以适配不同的架构
    (CPU, GPU 等)或不同的输入/输出类型.
  • 创建一个 Python 包装器（wrapper）. 这个包装器是创建 Op 的公开 API. 当注册 Op 时, 会自动生成一个默认
    默认的包装器. 既可以直接使用默认包装器, 也可以添加一个新的包装器.
  • (可选) 写一个函数计算 Op 的梯度.
  • (可选) 写一个函数, 描述 Op 的输入和输出 shape. 该函数能够允许从 Op 推断 shape.
  • 测试 Op, 通常使用 Pyhton。如果你定义了梯度，你可以使用Python的GradientChecker来测试它。

内容

增加一个新 Op

定义 Op 的接口

向 TensorFlow 系统注册来定义 Op 的接口. 在注册时, 指定 Op 的名称, 它的输入(类型和名称)
和输出(类型和名称), 和所需要任何 属性的文档说明.

为了让你有直观的认识, 创建一个简单的 Op 作为例子. 该 Op 接受一个 int32 类型 tensor 作为
输入, 输出这个 tensor 的一个副本, 副本与原 tensor 唯一的区别在于第一个元素被置为 0. 创建
文件 tensorflow/core/user_ops/zero_out.cc, 并调用 REGISTER_OP 宏来定义 Op 的接口.

  1.  #include "tensorflow/core/framework/op.h"
  2. REGISTER_OP("ZeroOut")
  3.     .Input("to_zero: int32")
  4.     .Output("zeroed: int32");

ZeroOut Op 接受 32 位整型的 tensor to_zero 作为输入, 输出 32 位整型的 tensor zeroed.

命名的注意事项: Op 的名称必须是为唯一的, 并使用驼峰命名法. 以下划线 _ 开始的名称保留为内部使用.

为 Op 实现 kernel

在定义接口之后, 提供一个或多个 Op 的实现. 为这些 kernel 的每一个创建一个对应的类, 继承
OpKernel, 覆盖 Compute 方法. Compute 方法提供一个类型为 OpKernelContext* 的参数 context, 用于访问一些有用的信息, 例如输入和输出的 tensor.

将 kernel 添加到刚才创建的文件中, kernel 看起来和下面的代码类似:

  1.  #include "tensorflow/core/framework/op_kernel.h"
  2. using namespace tensorflow;
  3. class ZeroOutOp : public OpKernel {
  4.  public:
  5.   explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}
  6.   void Compute(OpKernelContext* context) override {
  7.     // 获取输入 tensor.
  8.     const Tensor& input_tensor = context->input(0);
  9.     auto input = input_tensor.flat<int32>();
  10.    // 创建一个输出 tensor.
  11.     Tensor* output_tensor = NULL;
  12.     OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
  13.                                                      &output_tensor));
  14.     auto output = output_tensor->template flat<int32>();
  15.     // 设置 tensor 除第一个之外的元素均设为 0.
  16.     const int N = input.size();
  17.     for (int i = 1; i < N; i++) {
  18.       output(i) = 0;
  19.     }
  20.     // 尽可能地保留第一个元素的值.
  21.     if (N > 0) output(0) = input(0);
  22.   }
  23. };

实现 kernel 后, 将其注册到 TensorFlow 系统中. 注册时, 可以指定该 kernel 运行时的多个约束
条件. 例如可以指定一个 kernel 在 CPU 上运行, 另一个在 GPU 上运行.

将下列代码加入到 zero_out.cc 中, 注册 ZeroOut op:

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

一旦创建和重新安装了 TensorFlow ,
Tensorflow 系统可以在需要时引用和使用该 Op.

生成客户端包装器

Python Op 包装器

当编译 TensorFlow 时, 所有放在 tensorflow/core/user_ops 目录下
的 Op 会自动在 bazel-genfiles/tensorflow/python/ops/gen_user_ops.py 文件
中生成 Python Op 包装器. 通过以下声明, 把那些 Op 引入到 tensorflow/python/user_ops/user_ops.py
中:

  1. from tensorflow.python.ops.gen_user_ops import *

你可以选择性将部分函数替换为自己的实现. 为此, 首先要隐藏自动生成的代码,
在 tensorflow/python/BUILD
文件中, 将其名字添加到 "user_ops" 的 hidden 列表.

  1. tf_gen_op_wrapper_py(
  2.     name = "user_ops",
  3.     hidden = [
  4.         "Fact",
  5.     ],
  6.     require_shape_functions = False,
  7. )

紧接着 "Fact" 列出自己的 Op. 然后, 在
tensorflow/python/user_ops/user_ops.py
中添加你的替代实现函数. 通常, 替代实现函数也会调用自动生成函数来真正把 Op 添加
到图中. 被隐藏的自动生成函数位于 gen_user_ops 包中, 名称多了一个下划线前缀
(“_“). 例如:

  1. def my_fact():
  2.     """覆盖一个 Op 自动生成代码的示例."""
  3.     return gen_user_ops._fact()

C++ Op 包装器

当编译 TensorFlow 时, 所有 tensorflow/core/user_ops 文件夹
下的 Op 会自动创建 C++ Op 包装器. 例如, tensorflow/core/user_ops/zero_out.cc 中的 Op 会自动在 bazel-genfiles/tensorflow/cc/ops/user_ops.{h,cc}
中生成包装器.

tensorflow/cc/ops/standard_ops.h 通过下述申明,
导入用户自定义 Op 自动生成的包装器.

  1.  #include "tensorflow/cc/ops/user_ops.h"

检查 Op 能否正常工作

验证已经成功实现 Op 的方式是编写测试程序. 创建文件
tensorflow/python/kernel_tests/zero_out_op_test.py,
包含以下内容:

  1. import tensorflow as tf
  2. class ZeroOutTest(tf.test.TestCase):
  3.   def testZeroOut(self):
  4.     with self.test_session():
  5.       result = tf.user_ops.zero_out([5, 4, 3, 2, 1])
  6.       self.assertAllEqual(result.eval(), [5, 0, 0, 0, 0])

然后运行测试:

  1. $ bazel test tensorflow/python:zero_out_op_test

验证条件

上述示例假定 Op 能够应用在任何 shape 的 tensor 上. 如果只想应用到 vector 上
呢?
这意味需要在上述 OpKernel 实现中添加相关的检查.

  1.   void Compute(OpKernelContext* context) override {
  2.    // 获取输入 tensor
  3.     const Tensor& input_tensor = context->input(0);
  4.     OP_REQUIRES(context, TensorShapeUtils::IsVector(input_tensor.shape()),
  5.                 errors::InvalidArgument("ZeroOut expects a 1-D vector."));
  6.     // ...
  7.   }

OP_REQUIRES 断言的输入是一个 vector, 如果不是 vector, 将设置 InvalidArgument 状态并返回.
OP_REQUIRES 宏 有三个参数:

如果想要测试一个函数返回的 Status 对象是否是一个错误, 可以使用 OP_REQUIRES_OK.
这些宏如果检测到错误, 会直接跳出函数, 终止函数执行.

Op 注册

属性

Op 可以有属性, 属性的值在 Op 添加到图中时被设置. 属性值用于配置 Op, 在 kernel 实现中, Op 注册的输入和输出类型中, 均可访问这些属性值. 尽可能地使用输入代替属性, 因为输入的灵活性更高, 例如可以在执行步骤中
中被更改, 可以使用 feed 等等. 属性可用于实现一些输入无法做到的事情, 例如影响 Op 签名 (即输入输出的数量和类型)
的配置或只读配置可以通过属性实现.

注册 Op 时可以用 Attr 方法指定属性的名称和类型, 以此来定义一个属性, 形式如下:

  1. <name>: <attr-type-expr>

<name> 必须以字母开头, 可以由数字, 字母, 下划线组成. <attr-type-expr> 是一个类型表达式,
形式如下:

例如, 如果想要 ZeroOut Op 保存一个用户索引, 指示该 Op 不仅仅只有一个元素, 你可以注册 Op 如下:

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

你的 kernel 可以在构造函数里, 通过 context 参数访问这个属性:

  1. class ZeroOutOp : public OpKernel {
  2.  public:
  3.   explicit ZeroOutOp(OpKernelConstruction * context) : OpKernel(context) {
  4.    // 获取欲保存的索引值
  5.     OP_REQUIRES_OK(context,
  6.                    context->GetAttr("preserve_index", &preserve_index_));
  7.     // 检查 preserve_index 是否为正
  8.     OP_REQUIRES(context, preserve_index_ >= 0,
  9.                 errors::InvalidArgument("Need preserve_index >= 0, got ",
  10.                                         preserve_index_));
  11.   }
  12.   void Compute(OpKernelContext* context) override {
  13.     // ...
  14. }
  15.  private:
  16.   int preserve_index_;
  17. };

该值可以在 Compute 方法中被使用:

  1. void Compute(OpKernelContext* context) override {
  2.     // ...
  3.    // 检查 preserve_index 范围是否合法
  4. OP_REQUIRES(context, preserve_index_ < input.dimension(0),
  5.                 errors::InvalidArgument("preserve_index out of range"));
  6.     // 设置输出 tensor 所有的元素值为 0
  7.    const int N = input.size();
  8.     for (int i = 0; i < N; i++) {
  9.       output_flat(i) = 0;
  10.     }
  11.     // 保存请求的输入值
  12.    output_flat(preserve_index_) = input(preserve_index_);
  13.   }

为了维持向后兼容性, 将一个属性添加到一个已有的 Op 时,
必须指定一个默认值:

  1. REGISTER_OP("ZeroOut")
  2.      .Attr("preserve_index: int = 0")
  3.      .Input("to_zero: int32")
  4.      .Output("zeroed: int32");

属性类型

属性可以使用下面的类型:

  • string: 任何二进制字节流 (UTF8 不是必须的).
  • int: 一个有型整数.
  • float: 一个浮点数.
  • bool: 真或假.
  • type: DataType 非引用类型之一.
  • shape: 一个 TensorShapeProto.
  • tensor: 一个 TensorProto.
  • list(<type>): <type> 列表, 其中 <type> 是上述类型之一.
    注意 list(list(<type>)) 是无效的.

权威的列表以 op_def_builder.cc:FinalizeAttr 为准.

默认值和约束条件

属性可能有默认值, 一些类型的属性可以有约束条件. 为了定义一个有约束条件的属性, 你可以使用下列的
<attr-type-expr> 形式:

  • {'<string1>', '<string2>'}: 属性值必须是一个字符串, 取值可以为 <string1> 或 <string2>.
    值的语法已经暗示了值的类型为 string, 已经暗示了. 下述语句模拟了一个枚举值:
  1. REGISTER_OP("EnumExample")
  2.       .Attr("e: {'apple', 'orange'}");
  • {<type1>, <type2>}: 值是 type 类型, 且必须为 <type1> 或 <type2> 之一, 当然
    <type1> 和 <type2> 必须都是有效的 tensor 类型.
    你无须指定属性的类型为 type, 而是通过 {...} 语句给出一个类型列表. 例如, 在下面的例子里,
    属性 t 的类型必须为 int32, float, 或 bool:
  1. REGISTER_OP("RestrictedTypeExample")
  2.       .Attr("t: {int32, float, bool}");

  • 这里有一些常见类型约束条件的快捷方式:

    • numbertype: 限制类型为数字类型, 即非 string 非 bool 的类型.
    • realnumbertype: 与 numbertype 区别是不支持复杂类型.
    • quantizedtype: 与 numbertype 区别是只支持量化数值 (quantized number type).

这些类型的列表在 tensorflow/core/framework/types.h
文件中通过函数定义 (如 NumberTypes()).
本例中属性 t 必须为某种数字类型:

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

对于这个 Op:

  1. tf.number_type(t=tf.int32)  # 有效
  2. tf.number_type(t=tf.bool)   # 无效
  • int >= <n>: 值必须是一个整数, 且取值大于等于 <n>, <n> 是一个自然数.

例如, 下列 Op 注册操作指定了属性 a 的取值至少为 2.

  1. REGISTER_OP("MinIntExample")
  2.       .Attr("a: int >= 2");
  • list(<type>) >= <n>: 一个 <type> 类型列表, 列表长度必须大于等于 <n>.

例如, 下面的 Op 注册操作指定属性 a 是一个列表, 列表中的元素类型是 int32 或 float列表长度至少为3.

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

通过添加 = <default> 到约束条件末尾, 给一个属性设置默认值 (使其在自动生成的代码里
变成可选属性), 如下:

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

默认值支持的语法将在最终 GraphDef 定义的 protobuf 表示中被使用.

下面是给所有类型赋予默认值的例子:

  1. REGISTER_OP("AttrDefaultExampleForAllTypes")
  2.    .Attr("s: string = 'foo'")
  3.    .Attr("i: int = 0")
  4.    .Attr("f: float = 1.0")
  5.    .Attr("b: bool = true")
  6.    .Attr("ty: type = DT_INT32")
  7.    .Attr("sh: shape = { dim { size: 1 } dim { size: 2 } }")
  8.    .Attr("te: tensor = { dtype: DT_INT32 int_val: 5 }")
  9.    .Attr("l_empty: list(int) = []")
  10.    .Attr("l_int: list(int) = [2, 3, 5, 7]");

请特别注意那些类型值里面包含的 DT_* 名称.

多态

Type Polymorphism

对于那些可以使用不同类型输入或产生不同类型输出的 Op, 可以注册 Op 时为输入/输出类型里指定一个属性.
一般紧接着, 会为每一个支持的类型注册一个 OpKernel.

例如, 除了 int32 外, 想要 ZeroOut Op 支持 float, 注册代码如下:

  1. REGISTER_OP("ZeroOut")
  2.     .Attr("T: {float, int32}")
  3.     .Input("to_zero: <b>T</b>")
  4.     .Output("zeroed: <b>T</b>");

这段 Op 注册代码现在指定了输入的类型必须为 float 或 int32, 而且
既然输入和输出制定了同样的类型 T, 输出也同样如此.

一个命名建议:{#naming} 输入, 输出, 和属性通常使用 snake_case 命名法.
唯一的例外是属性被用作输入类型或是输入类型的一部分. 当添加到图中时, 这些属性
可以被推断出来, 因此不会出现在 Op 的函数里. 例如, 最后一个 ZeroOut 定义
生成的 Python 函数如下:

  1. def zero_out(to_zero, name=None):
  2.    """...
  3.    参数:
  4.      to_zero: 一个 `Tensor`. 必须为下列类型之一:
  5.          `float32`, `int32`.
  6.      name: 操作的名字 (可选).
  8.    返回值:
  9.      一个 `Tensor`, 类型和 `to_zero` 一样.
  10.    """

如果输入的 to_zero 是一个 int32 的tensor, 然后 T 将被自动
设置为 int32 (实际上是 DT_INT32). 那些推导出的属性的名称字母全大写
或采用驼峰命名法.

下面是一个输出类型自动推断的例子, 读者可以对比一下:

  1. REGISTER_OP("StringToNumber")
  2.      .Input("string_tensor: string")
  3.      .Output("output: out_type")
  4.      .Attr("out_type: {float, int32}");
  5.      .Doc(R"doc(
  6.  Converts each string in the input Tensor to the specified numeric type.
  7.  )doc");

在这种情况下, 用户需要在生成的 Python 代码中指定输出类型.

  1. def string_to_number(string_tensor, out_type=None, name=None):
  2.    """将输入 Tensor 中的每一个字符串转化成指定的数字类型
  4.    参数:
  5.      string_tensor: 一个 `string` 类型的 `Tensor`.
  6.      out_type: 一个可选的 `tf.DType`, 取值为 `tf.float32, tf.int32`.
  7.        默认值是 `tf.float32`.
  8.      name: 操作的名称 (可选).
  10.    返回值:
  11.      一个 `out_type` 类型的 `Tensor`.
  12.    """
  1.  #include "tensorflow/core/framework/op_kernel.h"
  2. class ZeroOutInt32Op : public OpKernel {
  3.   // 和之前一样
  4. };
  5. class ZeroOutFloatOp : public OpKernel {
  6.  public:
  7.   explicit ZeroOutFloatOp(OpKernelConstruction * context)
  8.       : OpKernel(context) {}
  9.   void Compute(OpKernelContext * context) override {
  10.     // 获取输入 tensor
  11.     const Tensor& input_tensor = context->input(0);
  12.     auto input = input_tensor.flat<float>();
  13.     // 创建一个输出 tensor
  14.     Tensor * output = NULL;
  15.     OP_REQUIRES_OK(context,
  16.                     context->allocate_output(0, input_tensor.shape(), &output));
  17.     auto output_flat = output->template flat<float>();
  18.     // 设置输出 tensor 的所有元素为 0
  19.     const int N = input.size();
  20.     for (int i = 0; i &lt; N; i++) {
  21.       output_flat(i) = 0;
  22.     }<br/>
  23.     // 保留第一个输入值
  24.     if (N &gt; 0) output_flat(0) = input(0);
  25.   }
  26. };
  27. // 注意, TypeConstraint<int32>("T") 意味着属性 "T" (在上面 Op 注册代码中
  28. // 定义的) 必须是 "int32", 才能实例化. 
  29. REGISTER_KERNEL_BUILDER(
  30.     Name("ZeroOut")
  31.     .Device(DEVICE_CPU)
  32.     .TypeConstraint&lt;int32&gt;("T"),
  33.     ZeroOutOpInt32);
  34. REGISTER_KERNEL_BUILDER(
  35.     Name("ZeroOut")
  36.     .Device(DEVICE_CPU)
  37.     .TypeConstraint<float>("T"),
  38.     ZeroOutFloatOp);

为了保持向后兼容性, 你在为一个
已有的 op 添加属性时, 必须指定一个默认值:

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

如果需要添加更多类型, 例如 double:

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

为了避免为新增的类型写冗余的 OpKernel 代码, 通常可以写一个 C++ 模板作为替代.
当然, 仍然需要为每一个重载版本定义一个 keneral 注册 (REGISTER\_KERNEL\_BUILDER 调用).

  1. template <typename T>;
  2. class ZeroOutOp : public OpKernel {
  3.  public:
  4.     explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}
  5.   void Compute(OpKernelContext* context) override {
  6.     // 获取输入 tensor
  7.      const Tensor& input_tensor = context->input(0);
  8.     auto input = input_tensor.flat<T>();
  9.     // 创建一个输出 tensor
  10.       Tensor* output = NULL;
  11.     OP_REQUIRES_OK(context,
  12.                    context->allocate_output(0, input_tensor.shape(), &output));
  13.     auto output_flat = output->template flat<T>();
  14.     // 设置输出 tensor 的所有元素为 0
  15.    const int N = input.size();
  16.     for (int i = 0; i < N; i++) {
  17.       output_flat(i) = 0;
  18.     }
  19.     // Preserve the first input value
  20.     if (N > 0) output_flat(0) = input(0);
  21.   }
  22. };
  23. };<br/>
  24. // 注意, TypeConstraint<int32>("T") 意味着属性 "T" (在上面 Op 注册代码中
  25. // 定义的) 必须是 "int32", 才能实例化. </b>
  26. REGISTER_KERNEL_BUILDER(
  27.     Name("ZeroOut")
  28.     .Device(DEVICE_CPU)
  29.     .TypeConstraint<int32>("T"),
  30.     ZeroOutOp<int32>);
  31. REGISTER_KERNEL_BUILDER(
  32.     Name("ZeroOut")
  33.     .Device(DEVICE_CPU)
  34.     .TypeConstraint<float>("T"),
  35.     ZeroOutOp<float>);
  36. REGISTER_KERNEL_BUILDER(
  37.     Name("ZeroOut")
  38.     .Device(DEVICE_CPU)
  39.     .TypeConstraint<double>("T"),
  40.     ZeroOutOp<double>);

如果有很多重载版本, 可以将注册操作通过一个宏来实现.

  1.  #include "tensorflow/core/framework/op_kernel.h"
  2.  #define REGISTER_KERNEL(type)                                       \
  3.   REGISTER_KERNEL_BUILDER(                                          \
  4.       Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
  5.       ZeroOutOp<type>)
  6. REGISTER_KERNEL(int32);
  7. REGISTER_KERNEL(float);
  8. REGISTER_KERNEL(double);
  9.  #undef REGISTER_KERNEL

取决于注册 kernel 使用哪些类型, 你可能可以使用tensorflow/core/framework/register_types.h
提供的宏:

  1.  #include "tensorflow/core/framework/op_kernel.h"
  2.  #include "tensorflow/core/framework/register_types.h"
  3. REGISTER_OP("ZeroOut")
  4.     .Attr("T: realnumbertype")
  5.     .Input("to_zero: T")
  6.     .Output("zeroed: T");
  7. template <typename T>
  8. class ZeroOutOp : public OpKernel { ... };
  9.  #define REGISTER_KERNEL(type)                                       \
  10.   REGISTER_KERNEL_BUILDER(                                          \
  11.       Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
  12.       ZeroOutOp<type>)
  13. TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);
  14.  #undef REGISTER_KERNEL

列表输入和输出

除了能够使用不同类型的 tensor 作为输入或输出, Op 还支持使用多个 tensor 作为输入或输出.

在接下来的例子里, 属性 T 存储了一个类型列表, 并同时作为输入 in 和输出 out 的类型.
输入和输出均为指定类型的 tensor 列表. 既然输入和输出的类型均为 T, 它们的 tensor 数量和类型
是一致的.

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

可以为列表中可存放的类型设置约束条件. 在下一个例子中, 输入是 float 和
double 类型的 tensor 列表. 例如, 这个 Op 可接受的
输入类型为 (float, double, float) 的数据, 且在此情况下, 输出类型同样
为 (float, double, float).

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

如果想要一个列表中的所有 tensor 是同一类型, 你需要写下列代码:

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

这段代码接受 int32 tensor 列表, 并用一个 int 属性 N
来指定列表的长度.

这也可用于类型推断. 在下一个例子中,
输入是一个 tensor 列表, 长度为 "N", 类型为 "T", 输出是单个 "T" 的 tensor:

  1. REGISTER_OP("SameListInputExample")
  2.     .Attr("N: int")
  3.     .Attr("T: type")
  4.     .Input("in: N * T")
  5.     .Output("out: T");

默认情况下, tensor 列表的最小长度为1. 这个约束条件可以通过
为指定的属性增加一个 ">=" 约束来变更:

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

同样的语法也适用于 "list(type)" 属性:

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

输入和输出

总结一下上述内容, 一个 Op 注册操作可以指定多个输入和输出:

  1. REGISTER_OP("MultipleInsAndOuts")
  2.     .Input("y: int32")
  3.     .Input("z: float")
  4.     .Output("a: string")
  5.     .Output("b: int32");

每一个输入或输出形式如下:

  1. <name>: <io-type-expr>

其中, <name> 以字母打头, 且只能由数字, 字母和下划线组成. <io-type-expr> 可以是
下列类型表达式之一:

  • <type>, 一个合法的输入类型, 如 float, int32, string. 这可用于指定给定类型的单个 tensor.

参见合法 Tensor 类型列表.

  1. REGISTER_OP("BuiltInTypesExample")
  2.       .Input("integers: int32")
  3.       .Input("complex_numbers: scomplex64");
  • <attr-type>, 一个属性和一个类型 type 或类型列表 list(type)(可能
    包含类型限制). 该语法可实现多态 Op.
  1. REGISTER_OP("PolymorphicSingleInput")
  2.       .Attr("T: type")
  3.       .Input("in: T);
  4. REGISTER_OP("RestrictedPolymorphicSingleInput")
  5.       .Attr("T: {int32, int64}")
  6.       .Input("in: T);

将属性的类型设置为 list(type) 将允许你接受一个序列的 tensor.

  1. REGISTER_OP("ArbitraryTensorSequenceExample")
  2.       .Attr("T: list(type)")
  3.       .Input("in: T")
  4.       .Output("out: T");
  5. REGISTER_OP("RestrictedTensorSequenceExample")
  6.       .Attr("T: list({int32, int64})")
  7.       .Input("in: T")
  8.       .Output("out: T");

注意, 输入和输出均为 T, 意味着输入和输出的类型与数量均相同.

  • <number> * <type>, 一组拥有相同类型的 tensor, <number> 是一个 int 类型属性的名称.
    <type> 可以是一个类似于 int32 和 float 的特定类型,
    或者一个 type 类型属性的名字. 前者的例子如下, 该例子接受一个 int32 tensor 列表作为 Op 输入:
  1. REGISTER_OP("Int32SequenceExample")
  2.       .Attr("NumTensors: int")
  3.       .Input("in: NumTensors * int32")

后者的例子如下, 该例子接受一个泛型 tensor 列表作为 Op 输入:

  1. REGISTER_OP("SameTypeSequenceExample")
  2.       .Attr("NumTensors: int")
  3.       .Attr("T: type")
  4.       .Input("in: NumTensors * T")
  • Tensor 的引用表示为 Ref(<type>), 其中 <type> 是上述类型之一.

一个命名建议: 当使用属性表示一个输入的类型时, 该类型可以被推断出来. 实现该特性, 将需要推断
的类型用大写名称表示 (如 T 或 N), 其它的输入, 输出, 和属性像使用函数参数一样使用这些
大写名称. 参见之前的命名建议章节查看更多细节.

更多细节参见 tensorflow/core/framework/op_def_builder.h.

向后兼容性

通常, 对规范的改变必须保持向后兼容性: Op 使用新规范后, 需保证使用旧规范构造的序列化 GraphDef 仍能正确工作.

下面是几种保持向后兼容性的方式:

  1. 任何添加到 Op 的新属性必须有默认值, 且默认值下的行为有明确定义. 将一个非多态的操作变为多态操作,
    你必须为新的类型属性赋予默认值, 以保持原始的函数签名. 例如, 有如下操作:
  1. REGISTER_OP("MyGeneralUnaryOp")
  2.        .Input("in: float")
  3.        .Output("out: float");

可以通过下述方式将其变为多态, 且保持向后兼容性:

  1. REGISTER_OP("MyGeneralUnaryOp")
  2.        .Input("in: T")
  3.        .Output("out: T")
  4.        .Attr("T: numerictype = float");

1.放宽一个属性的约束条件是安全的. 例如, 你可以将 {int32, int64} 变为 {int32, int64, float},
或者, 将 {"apple", "orange"} 变为 {"apple", "banana", "orange"}.

2.通过给 Op 名称添加一些项目中唯一的标识作为前缀, 来为新建的 Op 添加命名空间. 命名空间
可以预防你的 Op 与 TensorFlow 未来版本里的内置 Op 产生命名冲突.

3.超前计划! 尝试着去预测 Op 未来的的用途, 超前设计, 毕竟, 一些签名的变更无法保证兼容性
(例如, 增加新的输入, 或将原来的单元素输入变成一个列表).

如果不能以兼容的方式改变一个操作, 那就创建一个全新的操作, 来实现所需功能.

GPU 支持

你可以实现不同的 OpKernel, 将其中之一注册到 GPU, 另一个注册到 GPU, 正如为不同的类型注册 kernel 一样.
tensorflow/core/kernels/ 中有一些 GPU 支持的例子.
注意, 一些 kernel 的 CPU 版本位于 .cc 文件, GPU 版本位于 _gpu.cu.cc 文件, 共享的代码位于 .h 文件.

例如, pad op 除了 GPU kernel 外的其它代码
均在 tensorflow/core/kernels/pad_op.cc 中. GPU kernel 位于 tensorflow/core/kernels/pad_op_gpu.cu.cc,
共享的一个模板类代码定义在 tensorflow/core/kernels/pad_op.h.
需要注意的事情是, 即使使用 pad 的 GPU 版本时, 仍然需要将 "paddings" 输入放置到内存中.
为了实现这一点, 将输入或输出标记为必须保存在内存中, 为 kernel 注册一个 HostMemory() 调用.
如下:

  1.  #define REGISTER_GPU_KERNEL(T)                         \
  2. REGISTER_KERNEL_BUILDER(Name("Pad")                  \
  3.                               .Device(DEVICE_GPU)      \
  4.                               .TypeConstraint<T>("T")  \
  5.                               .HostMemory("paddings"), \
  6.                           PadOp<GPUDevice, T>)

使用 Python 实现梯度

给定一个 Op 组成的图, TensorFlow 使用自动微分 (反向传播) 来添加新的 Op 以表示梯度运算, 同时
不影响已有的 Op (参见梯度运算).
为了使自动微分能够与新的 Op 协同工作, 必须注册一个梯度函数, 从 Op 的输入计算梯度, 并返回代表
梯度值的输出.

数学上, 如果一个 Op 计算 \(y = f(x)\), 注册的梯度 Op 通过以下链式法则, 将 \(\partial / \partial y\)
的梯度运算转化为 \(\partial / \partial x\) 的梯度运算.

\frac{\partial}{\partial x}
= \frac{\partial}{\partial y} \frac{\partial y}{\partial x}
= \frac{\partial}{\partial y} \frac{\partial f}{\partial x}.

在 ZeroOut 的例子中, 输入中只有一个项会影响输出, 所以, 代表输入的梯度值的 tensor 也只有
一个输入项. 如下所示:

  1. from tensorflow.python.framework import ops
  2. from tensorflow.python.ops import array_ops
  3. from tensorflow.python.ops import sparse_ops
  5. @ops.RegisterGradient("ZeroOut")
  6. def _zero_out_grad(op, grad):
  7.   """`zero_out` 的梯度.
  9.   参数:
  10.     op: 欲进行微分的 `zero_out` `操作`, 可以用于获取原始 Op 的输入和输出.
  11.     grad: 代表 `zero_out` 输出的梯度 Op.
  13.   返回:
  14.     代表输入 `zero_out` 的微分.
  15.   """
  16.   to_zero = op.inputs[0]
  17.   shape = array_ops.shape(to_zero)
  18.   index = array_ops.zeros_like(shape)
  19.   first_grad = array_ops.reshape(grad, [-1])[0]
  20.   to_zero_grad = sparse_ops.sparse_to_dense(index, shape, first_grad, 0)
  21.   return [to_zero_grad]  # 单个 Tensor 的列表, 既然只有一个输入

使用 ops.RegisterGradient
注册梯度函数需要注意的一些细节:

  • 对于仅有一个输出的 Op, 梯度函数使用 Operation op
    和一个 Tensor grad 作为参数, 并从
    op.inputs[i],
    op.outputs[i],
    和 grad 构建新的 Op. 属性的信息可以通过 op.get_attr 获取.

  • 如果 Op 有多个输出, 梯度函数将使用 op 和 grads 作为参数, 其中, grads 是一个
    梯度 Op 的列表, 为每一个输出计算梯度. 梯度函数的输出必须是一个 Tensor 对象列表, 对应到
    每一个输入的梯度.

  • 如果没有为一些输入定义梯度, 譬如用作索引的整型, 这些输入返回的梯度为 None. 举一个例子,
    如果一个 Op 的输入为一个浮点数 tensor x 和一个整型索引 i, 那么梯度函数将返回
    [x_grad, None].

  • 如果梯度对于一个 Op 来说毫无意义, 使用 ops.NoGradient("OpName") 禁用自动差分.

注意当梯度函数被调用时, 作用的对象是数据流图中的 Op, 而不是 tensor 数据本身. 因此, 只有在图运行时,
梯度运算才会被其它 tensorflow Op 的执行动作所触发.

在 Python 中实现一个形状函数

TensorFlow Python API 有一个 “形状推断” 功能, 可以不执行图就获取 tensor 的形状信息.
形状推断功能藉由每一个 Op 类型注册的 “形状函数” 来支持, 该函数有两个规则: 假设所有输入的
形状必须是兼容的, 以及指定输出的形状. 一个形状函数以一个 Operation
作为输入, 返回一个 TensorShape
对象列表 (每一个输出一个对象). 使用 tf.RegisterShape 装饰器
注册形状函数. 例如, 上文定义的 ZeroOut Op 的形状函数如下:

  1. @tf.RegisterShape("ZeroOut"):
  2. def _zero_out_shape(op):
  3.   """ZeroOut Op 的形状函数.
  5.   这是 ZeroOut 形状函数的无约束版本, 为每一个输出产生的形状和对应的输入一样. 
  6.   """
  7.   return [op.inputs[0].get_shape()]

一个形状函数也可以约束输入的形状. 下面是 ZeroOut 形状函数的 vector 输入约束版本:

  1. @tf.RegisterShape("ZeroOut"):
  2. def _zero_out_shape(op):
  3.   """ZeroOut Op 的形状函数.
  5.   这是 ZeroOut 形状函数的约束版本, 要输入的 rank 必须是 1 (即使一个 vector).
  6.   """
  7.   input_shape = op.inputs[0].get_shape().with_rank(1)
  8.   return [input_shape]

如果 Op 是多输入的多态 Op, 使用操作的属性来决定需要检查的形状数量:

  1. @tf.RegisterShape("IntListInputExample")
  2. def _int_list_input_example_shape(op):
  3.   """ "IntListInputExample" Op 的形状函数.
  5.   所有的输入和输出是同大小的矩阵.
  6.   """
  7.   output_shape = tf.TensorShape(None)
  8.   for input in op.inputs:
  9.     output_shape = output_shape.merge_with(input.get_shape().with_rank(2))
  10.   return [output_shape]

既然形状推断是一个可选的特性, 且 tensor 的形状可能动态变化, 形状函数必须足够健壮, 能够处理任意
输入形状信息缺失的情形. merge_with 方法能够帮助
调用者判断两个形状是否是一样的, 即使两个形状的信息不全, 该函数同样有效. 所有的标准 Python Op
的形状函数都已经定义好了, 并且已经有很多不同的使用示例.

原文：Adding a New Op 翻译：@doc001 校对：@ZHNathanielLee