总结一下上述内容, 一个 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"); |
<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") |
| REGISTER_OP("SameTypeSequenceExample") |
| .Attr("NumTensors: int") |
| .Attr("T: type") |
| .Input("in: NumTensors * T") |
Ref(<type>), 其中 <type> 是上述类型之一.
通常, 对规范的改变必须保持向后兼容性: Op 使用新规范后, 需保证使用旧规范构造的序列化 GraphDef 仍能正确工作.
下面是几种保持向后兼容性的方式:
| 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 未来的的用途, 超前设计, 毕竟, 一些签名的变更无法保证兼容性 (例如, 增加新的输入, 或将原来的单元素输入变成一个列表).
如果不能以兼容的方式改变一个操作, 那就创建一个全新的操作, 来实现所需功能.
你可以实现不同的 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>) |
给定一个 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].
ops.NoGradient("OpName") 禁用自动差分.
注意当梯度函数被调用时, 作用的对象是数据流图中的 Op, 而不是 tensor 数据本身. 因此, 只有在图运行时, 梯度运算才会被其它 tensorflow Op 的执行动作所触发.
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()] |
| @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] |
