tensorflow 转张量类型为float_TensorFlow 指南：添加新操作（二）

将高级功能编译到操作中

昨天我们发布了“TensorFlow 指南：添加新操作(一)”，阅读过的朋友，您应该已经知道如何编译一个基本(并受到一些限制的)操作和实现，今天我们来看看您通常需要编译到操作中的一些更复杂的内容。

条件检查和验证

上面的示例假设操作已应用于任意形状的张量。如果操作只应用于向量，该怎么办？这就需要在上面的 OpKernel 实现中添加检查。

  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."));
    // ...
  }

上述代码会断言输入是一个向量，如果不是，则返回设置 InvalidArgument 状态。OP_REQUIRES 宏采用三个参数：

• context，可以是其 SetStatus() 方法的 OpKernelContext 或 OpKernelConstruction 指针(请参阅 tensorflow/core/framework/op_kernel.h https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h)

• 条件。例如，tensorflow/core/framework/tensor_shape.h 中有用于验证张量形状的函数

• 错误本身，由 Status 对象表示，请参阅 tensorflow/core/lib/core/status.h。Status 包含类型(通常为 InvalidArgument，但具体请参见类型列表)和消息。可以在 tensorflow/core/lib/core/errors.h 中找到用于构造错误的函数

或者，如果您要测试从某个函数返回的 Status 对象是否有误，请使用 OP_REQUIRES_OK(如果是，就返回它)。这两个宏都会在出错时从函数返回(https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/lib/core/errors.h)。

操作注册

属性

操作可以包含属性，其值是在向图中添加操作时设置的。这些值用于配置操作，并且您可以在内核实现中以及操作注册中的输入和输出类型中访问它们的值。尽可能首选使用输入而不是属性，因为输入更灵活。这是因为属性是常量，必须在构建图时定义。相比之下，输入是张量，其值可以动态变化；也就是说，输入在每一步都可以更改，通过馈送数据进行设置等。属性用于无法通过输入设置的方面：任何影响函数签名(输入或输出的数量或类型)或者不能一步一步进行更改的配置。

您可以在注册操作时定义属性，只需使用 Attr 方法(要求满足以下格式规范)指定该属性的名称和类型即可：

:

其中  以字母开头，可以由字母数字字符与下划线组成， 是如下所述形式的类型表达式。

例如，如果您希望 ZeroOut 操作保留用户指定的索引，而不是仅保留第 0 个元素，则可以如下所示地注册该操作：

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

(请注意，属性类型集不同于用于输入和输出的 tf.DType。)

然后，内核可以通过 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_;
};

然后，您可以在 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_);
  }

属性类型

属性支持以下类型：

• string：任意字节序列(无需是 UTF8)。

• int：有符号整数

• float：浮点数

• bool：True 或 False

• type：DataType 的(非引用)值之一

• shape：TensorShapeProto

• tensor：TensorProto

• list()： 列表，其中  是上述类型之一。请注意，list(list()) 无效

另请参阅 op_def_builder.cc:FinalizeAttr 以查看最终列表(https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_def_builder.cc)。

默认值和限制条件

属性可能具有默认值，并且某些类型的属性可能具有限制条件。要定义具有限制条件的属性，您可以使用以下 ：

• {'', ''}：值必须是值为  或  的字符串。使用此语法时，类型名称 string 隐含其中。下面模拟了一个枚举：

REGISTER_OP("EnumExample")
    .Attr("e: {'apple', 'orange'}");

• {, }：值的类型为 type，必须是  或  之一，其中  和  是受支持的 tf.DType。您未指定属性类型是 type。如果 {...} 中包含类型列表，该类型会隐含其中。例如，在以下示例中，属性 t 是一个必须是 int32、float 或 bool 的类型：

REGISTER_OP("RestrictedTypeExample")
    .Attr("t: {int32, float, bool}");

• 常见类型限制条件的快捷方式如下：

  numbertype：类型 type 仅限于数值(非字符串和非布尔)类型

  realnumbertype：类似于没有复杂类型的 numbertype

  quantizedtype：类似于 numbertype，但只是量化的数值类型

这些方式允许的特定类型列表由 tensorflow/core/framework/types.h 中的函数(如 NumberTypes())定义。在以下示例中，属性 t 必须是某个数值类型：

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

对于以下操作：

tf.number_type(t=tf.int32)  # Valid
tf.number_type(t=tf.bool)   # Invalid

列表可以与其他列表和单个类型组合。以下操作允许属性 t 为任意数值类型或布尔类型：

REGISTER_OP("NumberOrBooleanType")
    .Attr("t: {numbertype, bool}");

对于以下操作：

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 >= ：值必须是整型，其值大于或等于 ，其中  是自然数。

  例如，以下操作注册指定属性 a 的值必须至少为 2：

REGISTER_OP("MinIntExample")
    .Attr("a: int >= 2");

• list() >= ：类型为  的列表，其长度大于或等于 。

  例如，以下操作注册指定属性 a 是 int32 或 float 类型的列表，并且必须至少有 3 个值：

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

要设置属性的默认值(使其在生成的代码中可选)，请将 =  添加到末尾，如：

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

默认值支持的语法是在生成的 GraphDef 定义的 proto 表示法中使用的语法。

以下示例演示了如何为所有类型指定默认值：

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

特别要注意的是，类型为 type 的值要使用 tf.DType。

多态性

多态类型

对于可以采用不同类型的输入或生成不同类型的输出的操作，您可以在操作注册的输入或输出类型中指定属性。通常，您会为每种受支持的类型注册一个 OpKernel。

例如，如果除了 int32 之外您还想对 float 执行 ZeroOut 操作，那么操作注册可能如下所示：

REGISTER_OP("ZeroOut")
    .Attr("T: {float, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

操作注册现在指定输入类型必须是 float 或 int32，并且输出将是同一个类型，因为它们都具有 T 类型。

有关命名的注意事项：输入、输出和属性通常应该采用蛇形名称。但是，用作输入类型或用在输入类型中的属性例外。向图中添加操作时，系统可以推断出这些属性，因此这些属性不会显示在操作的函数中。例如，最后一个 ZeroOut 定义将生成如下所示的 Python 函数：

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`.
  """

如果向 to_zero 传递 int32 张量，则 T 自动设置为 int32(实际上是 DT_INT32)。这些推断属性会采用大写或驼峰式名称。

将此操作与具有确定输出类型的类型属性的操作进行比较：

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

在以下示例中，用户必须指定输出类型，如生成的 Python 中所示：

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`.
  """

#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();
    // 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();
    // 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("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("T"),
    ZeroOutOpInt32);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint("T"),
    ZeroOutFloatOp);

要保留向后兼容性，您应在将属性添加到现有操作时指定默认值：

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

假设您想添加更多类型，例如 double：

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

通常，您可以使用 C++ 模板，而不使用上面的冗余代码编写另一个 OpKernel。对于每个过载项，您仍然有一个内核注册(REGISTER_KERNEL_BUILDER 调用)。

template
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();
    // 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();
    // 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("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("T"),
    ZeroOutOp);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint("T"),
    ZeroOutOp);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint("T"),
    ZeroOutOp);

如果您有多个过载项，则可以将注册放在宏中。

#include "tensorflow/core/framework/op_kernel.h"
#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint("T"), \
      ZeroOutOp)
REGISTER_KERNEL(int32);
REGISTER_KERNEL(float);
REGISTER_KERNEL(double);
#undef REGISTER_KERNEL

根据您为其注册内核的类型列表，您可以使用 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
class ZeroOutOp : public OpKernel { ... };
#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint("T"), \
      ZeroOutOp)
TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);
#undef REGISTER_KERNEL

列出输入和输出

除了能够接受或生成不同类型之外，操作还可以使用或生成数量不定的张量。

在下一个示例中，属性 T 存储的是一个类型列表，并同时用作输入 in 和输出 out 的类型。输入和输出是该类型的张量列表(输出中的张量数量和类型与输入的相同，因为二者的类型都为 T)。

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

您还可以对列表中可指定的类型施加限制。在下一个示例中，输入是 float 和 double 张量的列表。例如，操作接受输入类型 (float, double, float)，在这种情况下，输出类型也将是 (float, double, float)。

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

如果您希望列表中的所有张量都具有相同类型，则可以运行如下命令：

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

此操作接受 int32 张量列表，并使用 int 属性 N 指定列表的长度。

类型也可以是多态类型。在下一个示例中，输入是同一(但未指定)类型 ("T") 的张量(长度为 "N")列表，输出是一个具有匹配类型的张量：

REGISTER_OP("SameListInputExample")
    .Attr("N: int")
    .Attr("T: type")
    .Input("in: N * T")
    .Output("out: T");

默认情况下，张量列表的最小长度为 1。您可以对相应属性设定 ">=" 限制条件，更改该默认值。在下一个示例中，输入是至少有 2 个 int32 张量的列表：

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

相同的语法也适用于 "list(type)" 属性：

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

输入和输出

总之，操作注册可以有多个输入和输出：

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

每个输入或输出规范的格式如下：

:

其中  以字母开头，并且可以包含字母数字字符与下划线。 是以下类型表达式之一：

• ，其中  是受支持的输入类型(例如 float、int32、string)。它会指定一个具有给定类型的张量。

  请参阅 tf.DType(https://tensorflow.google.cn/api_docs/python/tf/dtypes/DType?hl=zh-CN)

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

• ，其中  是类型为 type 或 list(type)(可能有类型限制)的属性的名称。此语法允许多态操作。

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

引用类型为 list(type) 的属性可让您接受张量序列。

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

请注意，输出 out 中的张量数量和类型与输入 in 中的张量数量和类型相同，因为二者的类型均为 T。

• 对于具有同一类型的张量序列： * ，其中  是类型为 int 的属性的名称。 可以是 tf.DType，也可以是类型为 type 的属性的名称。对于第一种情况，举个例子，此操作会接受 int32 张量列表：

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

此操作会接受任何类型的张量列表，只要它们类型都相同就可以：

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

• 对于张量引用：Ref()，其中  是之前的某个类型。

有关命名的注意事项：系统将推断出采用输入类型的任何属性。按照惯例，这些推断属性会使用大写名称(如 T 或 N)，否则输入、输出和属性会具有与函数参数类似的名称(例如 num_outputs)。如需了解详情，请参阅前面有关命名的注意事项。

如需了解详情，请参阅 

tensorflow/core/framework/op_def_builder.h(https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_def_builder.h)。

向后兼容性

假设您已经编写了一个很好的自定义操作并与其他人分享该操作，用户在使用您的操作时感到满意。不过，您想以某种方式更改操作。

通常，对现有检入规范的更改必须向后兼容：更改操作规范不得破坏之前根据旧规范构造的序列化 GraphDef 协议缓冲区。

您可以通过以下几种方式保持向后兼容性。

1. 添加到操作的所有新属性都必须定义默认值，并且在使用该默认值时，操作必须具有原始行为。要将操作从非多态更改为多态，您必须为新类型属性提供默认值，以便在默认情况下保留原始签名。例如，如果您的操作是：

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

   您可以使用以下方式以向后兼容的方式使该操作变为多态：

   REGISTER_OP("MyGeneralUnaryOp") .Input("in: T") .Output("out: T") .Attr("T: numerictype = DT_FLOAT")

2. 您可以安全地放宽对属性的限制。例如，您可以从 {int32, int64} 更改为 {int32, int64, float} 或 type，也可以从 {"apple", "orange"} 更改为 {"apple", "banana", "orange"} 或 string

3. 只要列表类型的默认值与旧签名匹配，您就可以将单个输入 / 输出更改为列表输入 / 输出

4. 您可以添加新的列表输入 / 输出(如果它默认为空)

5. 为您创建的所有新操作设定命名空间，方法是在操作名称前添加项目所独有的内容作为前缀。这样可以避免您的操作与可能包含在 TensorFlow 的未来版本中的任何操作发生冲突

6. 未雨绸缪！尝试预测操作的未来用途。某些签名更改无法以兼容的方式完成(例如，将相同类型的列表变成不同类型的列表)

可以在 

tensorflow/core/framework/op_compatibility_test.cc 中找到安全和不安全更改的完整列表。如果您无法向后兼容地更改操作，则创建新的操作，并用新语义设定新的名称(https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_compatibility_test.cc)。

另请注意，虽然这些更改可以保持 GraphDef 兼容性，但生成的 Python 代码可能会以与旧调用者不兼容的方式发生更改。要使 Python API 保持兼容，可以在手写 Python 封装容器中谨慎地进行更改，并保留旧签名(可能要在末尾添加新可选参数的情况除外)。通常，只有在 TensorFlow 更改主要版本时才可以进行不兼容的更改，并且这些更改必须符合 GraphDef 版本语义。

GPU 支持

您可以实现不同的 OpKernel 并为 CPU 和 GPU 各注册一个内核，就像您可以为不同类型注册内核一样。在 tensorflow/core/kernels/ 中有几个支持 GPU 的内核示例。请注意，某些内核在 .cc 文件中具有 CPU 版本，在以 _gpu.cu.cc 结尾的文件中具有 GPU 版本，并且在 .h 文件中具有一些共用代码。

例如，tf.pad 在 tensorflow/core/kernels/pad_op.cc 中具有除 GPU 内核之外的所有代码。GPU 内核位于 tensorflow/core/kernels/pad_op_gpu.cu.cc 中，共享代码是在 tensorflow/core/kernels/pad_op.h 中定义的模板化类。我们以这种方式组织代码有两个原因：可以让您在 CPU 和 GPU 实现之间共用代码，并且将 GPU 实现放入单独的文件中，以便它只能由 GPU 编译器编译。

有一点需要注意，即使使用 pad 的 GPU 内核版本，它在 CPU 内存中仍然需要 "paddings" 输入。要标记将输入或输出保留在 CPU 上，请添加对内核注册的 HostMemory() 调用，例如：

#define REGISTER_GPU_KERNEL(T)                         \
  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
                              .Device(DEVICE_GPU)      \
                              .TypeConstraint("T")  \
                              .HostMemory("paddings"), \
                          PadOp)

编译 GPU 设备的内核

要了解使用 CUDA 内核实现操作的示例，请参阅 

cuda_op_kernel.cu.cc(https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/adding_an_op/cuda_op_kernel.cu.cc)。

tf_custom_op_library 会接受 gpu_srcs 参数，可以在其中指定包含 CUDA 内核的源文件(*.cu.cc 文件)列表。要用于 TensorFlow 的二进制安装，必须使用 NVIDIA 的 nvcc 编译器编译 CUDA 内核。您可以使用以下命令序列将 cuda_op_kernel.cu.cc 和 cuda_op_kernel.cc 编译成一个可动态加载的库：

nvcc -std=c++11 -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++11 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \
  cuda_op_kernel.cu.o ${TF_CFLAGS[@]} -fPIC -lcudart ${TF_LFLAGS[@]}

可以像往常一样，在 Python 中使用 tf.load_op_library 函数加载上面生成的 cuda_op_kernel.so。

请注意，如果 CUDA 库未安装在 /usr/local/lib64 中，您需要在上面的第二个 (g++) 命令中明确指定路径。例如，如果 CUDA 安装在 /usr/local/cuda-8.0 中，请添加 -L /usr/local/cuda-8.0/lib64/。

请注意，在某些 Linux 设置中，需要在 nvcc 编译步骤中添加其他选项。将 -D_MWAITXINTRIN_H_INCLUDED 添加到 nvcc 命令行可以避免 mwaitxintrin.h 出错。

在 Python 中实现梯度

给定操作图后，TensorFlow 会使用自动微分(反向传播)添加表示梯度(相对于现有操作，请参见梯度计算)的新操作。要对新操作应用自动微分，必须注册一个梯度函数，以在给定相对于操作输出的梯度时，计算相对于操作输入的梯度。

在数学上，如果操作计算 \(y = f(x)\)，注册的梯度操作会通过链式法则将相对于 \(y\) 的损失 \(L\) 梯度 \(\partial L/ \partial y\) 转换为相对于 \(x\) 的梯度 \(\partial L/ \partial x\)：

$$\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}.$$

对于 ZeroOut，输入中只有一个条目会影响输出，因此相对于输入的梯度是稀疏的 “独热” 张量。具体表示方式如下：

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

有关使用 tf.RegisterGradient 注册梯度函数的详细信息如下所示：

• 对于具有一项输出的操作，梯度函数将采用 tf.Operation op 和 tf.Tensor grad，并根据张量 op.inputs[i]、op.outputs[i] 和 grad 构建新操作。所有属性的相关信息均可通过 tf.Operation.get_attr 找到

• 如果操作有多项输出，则梯度函数将采用 op 和 grads，其中 grads 是相对于每项输出的梯度列表。梯度函数的结果必须是 Tensor 对象(表示相对于每项输入的梯度)列表

• 如果某项输入(例如用作索引的整数输入)没有明确定义的梯度，则返回的相应梯度应为 None。例如，对于采用浮点张量 x 和整数索引 i 的操作，梯度函数将为 return [x_grad, None]

• 如果操作根本没有有意义的梯度，您通常无需注册任何梯度，并且只要从不需要该操作的梯度就可以。在某些情况下，操作没有明确定义的梯度，但可以参与梯度计算。在这种情况下，您可以使用 ops.NotDifferentiable 自动向后传播零

请注意，在调用梯度函数时，只有操作的数据流图可用，而张量数据本身不可用。因此，必须使用其他 TensorFlow 操作执行所有计算，以在图执行时运行。

C++ 中的形状函数

TensorFlow API 具有一项称为 “形状推断” 的功能，该功能可以提供有关张量形状的信息，而无需执行图。形状推断由在 C++ REGISTER_OP 声明中为每个操作类型注册的 “形状函数” 支持，并承担两个角色：声明输入的形状在图构造期间是兼容的，并指定输出的形状。

形状函数定义为对 shape_inference::InferenceContext 类的操作。例如，在 ZeroOut 的形状函数中：

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

c->set_output(0, c->input(0)); 会声明第一个输出的形状应该设置为第一个输入的形状。如果按输出索引选择输出(如上例所示)，则 set_output 的第二个参数应该是 ShapeHandle 对象。您可以通过空 ShapeHandle 对象的默认构造函数创建该对象。具有索引 idx 的输入的 ShapeHandle 对象可以通过 c->input(idx) 获得。

有许多常见的形状函数适用于多种操作(例如 shape_inference::UnchangedShape)，您可以在 common_shape_fns.h 中找到这些函数并按如下方式使用它们：

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);

形状函数还可以限制输入的形状。对于具有向量形状限制条件的 ZeroOut 版本，形状函数将如下所示：

    .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();
    });

WithRank 调用会验证输入形状 c->input(0) 是否具有只有一维的形状(或者，如果输入形状未知，则输出形状将为具有一个未知维度的向量)。

如果您的操作是多态的且包含多项输入，您可以使用 InferenceContext 成员确定要检查的形状的数量，并使用 Merge 验证这些形状是否均兼容(或者，使用可提供对操作属性的访问权限的 InferenceContext::GetAttr 来访问表示长度的属性)。

    .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();
    });

由于形状推断是可选功能，并且张量的形状可能会动态变化，因此对于所有输入的不完整形状信息，形状函数必须稳定可靠。InferenceContext 的 Merge 方法允许调用者声明两个形状是相同的，即使其中任一个或两个都没有完整的信息。我们为所有核心 TensorFlow 操作定义了形状函数，并提供了许多不同的用法示例。

InferenceContext 类具有许多可用于定义形状函数操作的函数。例如，您可以使用 InferenceContext::Dim 和 InferenceContext::WithValue 验证特定维度是否具有非常具体的值；您可以使用 InferenceContext::Add 和 InferenceContext::Multiply 指定输出维度是两个输入维度的和 / 积。要了解您可以指定的所有不同形状操作，请参阅 InferenceContext 类。以下示例将第一项输出的形状设置为 (n, 3)，其中第一项输入的形状为 (n, ...)

.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    c->set_output(0, c->Matrix(c->Dim(c->input(0), 0), 3));
    return Status::OK();
});

如果您具有复杂的形状函数，则应考虑添加测试，以验证各种输入形状组合是否会生成预期的输出形状组合。您可以在我们的一些核心操作测试中看到有关如何编写这些测试的示例。(INFER_OK 和 INFER_ERROR 的语法有点晦涩难懂，但在表示测试中的输入和输出形状规范时力求做到紧凑。暂时请查看这些测试中的注释，以了解形状字符串规范。)

更多 AI 相关阅读：

• 使用注意力机制制作图像字幕

• 欢迎试用 tf.function 加速代码

• TensorFlow 2.0 中的 tf.function 和 AutoGraph