输入和输出

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

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

<name>: <io-type-expr>

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

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

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

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

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

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

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

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

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

  Tensor 的引用表示为 Ref(<type>), 其中 <type> 是上述类型之一.

• 向后兼容性

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

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

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

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

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

  REGISTER_OP("MyGeneralUnaryOp")
         .Input("in: T")
         .Output("out: T")
         .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() 调用. 如下:

  #define REGISTER_GPU_KERNEL(T)                         \
  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
                                .Device(DEVICE_GPU)      \
                                .TypeConstraint<T>("T")  \
                                .HostMemory("paddings"), \
                            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 也只有 一个输入项. 如下所示:

  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):
    """`zero_out` 的梯度.
  
    参数:
      op: 欲进行微分的 `zero_out` `操作`, 可以用于获取原始 Op 的输入和输出.
      grad: 代表 `zero_out` 输出的梯度 Op.
  
    返回:
      代表输入 `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]  # 单个 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 的形状函数如下:

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

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

@tf.RegisterShape("ZeroOut"):
def _zero_out_shape(op):
  """ZeroOut Op 的形状函数.

  这是 ZeroOut 形状函数的约束版本, 要输入的 rank 必须是 1 (即使一个 vector).
  """
  input_shape = op.inputs[0].get_shape().with_rank(1)
  return [input_shape]

@tf.RegisterShape("IntListInputExample")
def _int_list_input_example_shape(op):
  """ "IntListInputExample" Op 的形状函数.

  所有的输入和输出是同大小的矩阵.
  """
  output_shape = tf.TensorShape(None)
  for input in op.inputs:
    output_shape = output_shape.merge_with(input.get_shape().with_rank(2))
  return [output_shape]