增加一个新 Op(三)

输入和输出

总结一下上述内容, 一个 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>, 一个合法的输入类型, 如 floatint32string. 这可用于指定给定类型的单个 tensor.
REGISTER_OP("BuiltInTypesExample")
      .Input("integers: int32")
      .Input("complex_numbers: scomplex64");
<attr-type>, 一个属性和一个类型 type 或类型列表 list(type)(可能 包含类型限制). 该语法可实现多态 Op.
REGISTER_OP("PolymorphicSingleInput")
      .Attr("T: type")
      .Input("in: T);
REGISTER_OP("RestrictedPolymorphicSingleInput")
      .Attr("T: {int32, int64}")
      .Input("in: T);
将属性的类型设置为 list(type) 将允许你接受一个序列的 tensor.
REGISTER_OP("ArbitraryTensorSequenceExample")
      .Attr("T: list(type)")
      .Input("in: T")
      .Output("out: T");
REGISTER_OP("RestrictedTensorSequenceExample")
      .Attr("T: list({int32, int64})")
      .Input("in: T")
      .Output("out: T");
注意, 输入和输出均为 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]
如果 Op 是多输入的多态 Op, 使用操作的属性来决定需要检查的形状数量:

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

联系我们

邮箱 626512443@qq.com
电话 18611320371
QQ 626512443
QQ群 235681453

Copyright © 2015-2019

备案号:京ICP备15003423号-3