TensorFlow 运算融合

概览

本页将介绍将 TensorFlow 中的复合操作转换为 TensorFlow Lite 中的融合操作所需的设计和步骤。此基础架构是通用的，并且支持将 TensorFlow 中的任何复合操作转换为 TensorFlow Lite 中的相应融合操作。

此基础架构的一个应用示例是将 TensorFlow RNN 操作融合到 TensorFlow Lite，详见此处。

什么是融合操作

绘图

TensorFlow 操作可以是基元操作（例如 tf.add），也可以由其他基元操作（例如 tf.einsum）组合而成。原始操作在 TensorFlow 图中显示为单个节点，而复合操作是 TensorFlow 图中的节点集合。执行复合操作等同于执行它的各个原初操作。

一个融合操作对应于单个操作，该操作包含相应复合操作中每个原始操作执行的所有计算。

融合操作的优势

融合操作旨在通过优化整体计算并减少内存占用量来最大限度地提高其底层内核实现的性能。这非常有用，尤其是对于低延迟推理工作负载和资源受限的移动平台。

一体化运算还提供更高级别的接口，用于定义量化等复杂转换，这在更精细的级别是不可行或非常难以执行的。

由于上述原因，TensorFlow Lite 有许多融合操作实例。这些融合操作通常对应于源 TensorFlow 程序中的复合操作。在 TensorFlow Lite 中作为单个融合操作实现的复合运算示例包括各种 RNN 运算，例如单向和双向序列 LSTM、卷积（conv2d、权重添加、relu）、全连接（matmul、bi 向添加、relu）等。在 TensorFlow Lite 中，LSTM 量化目前仅在融合 LSTM 操作中实现。

融合操作面临的挑战

在 TensorFlow Lite 中，将复合运算从 TensorFlow 转换为融合运算是一个难题。原因如下：

  1. 复合操作在 TensorFlow 图中表示为一组没有明确定义的边界的原始操作。识别（例如通过模式匹配）与此类复合运算对应的子图可能非常具有挑战性。

  2. 针对融合 TensorFlow Lite 操作的 TensorFlow 实现可能不止一个。例如，TensorFlow 中有许多 LSTM 实现（Keras、Babelfish/lingvo 等），每个实现都由不同的原语操作组成，但它们仍然可以在 TensorFlow Lite 中转换为相同的融合 LSTM 操作。

因此，事实证明，转换一体化操作的转换颇具挑战性。

将复合操作封装在 tf.function 中

在许多情况下，模型的某些部分可以映射到 TFLite 中的单个操作。这有助于提高针对特定操作编写优化实现时的性能。为了能够在 TFLite 中创建融合操作，请找出图中表示融合操作的部分，并将其封装在带有“experimental_implements”属性的 tf.function 中，并将其属性设为 tf.function（属性值为 tfl_fusable_op，值为 true）。如果自定义操作采用属性，则将它们作为同一“experimental_implements”的一部分传递。

例如，

def get_implements_signature():
  implements_signature = [
    # 'name' will be used as a name for the operation.
    'name: "my_custom_fused_op"',
    # attr "tfl_fusable_op" is required to be set with true value.
    'attr {key: "tfl_fusable_op" value { b: true } }',
    # Example attribute "example_option" that the op accepts.
    'attr {key: "example_option" value { i: %d } }' % 10
  ]
  return ' '.join(implements_signature)

@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
  # An empty function that represents pre/post processing example that
  # is not represented as part of the Tensorflow graph.
  output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
  output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
  return output_1, output_2

class TestModel(tf.Module):
  def __init__(self):
    super(TestModel, self).__init__()
    self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
    self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
  ])
  def simple_eval(self, input_a, input_b):
    return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))

请注意，您无需在转换器上设置 allow_custom_ops，因为 tfl_fusable_op 属性已暗示这一点。

实现自定义操作并注册 TFLite 解释器

将一体化操作实现为 TFLite 自定义操作 - 请参阅instructions。

请注意，用于注册操作的名称应与实现签名的 name 属性中指定的名称相似。

此示例中的操作示例如下：

  TfLiteRegistration reg = {};
  // This name must match the name specified in the implements signature.
  static constexpr char kOpName[] = "my_custom_fused_op";
  reg.custom_name = kOpName;
  reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.builtin_code = kTfLiteCustom;
  resolver->AddCustom(kOpName, &reg);

从复合操作转换为融合操作（高级）

将 TensorFlow 复合操作转换为 TensorFlow Lite 融合操作的总体架构如下所示：

绘图

将复合操作封装在 tf.function 中

在 TensorFlow 模型源代码中，找到复合操作并使用 experimental_implements 函数注解将其抽象化为 tf.function。请参阅嵌入查询的示例。此函数定义了接口及其参数，用于实现转换逻辑。

编写转化代码

转化代码按带 implements 注解的函数接口编写。请参阅嵌入查询的融合示例。从概念上讲，转换代码会将此接口的复合实现替换为融合实现。

在 prepare-composite-functions 传递后，在转换代码中添加插件。

在更高级的用法中，可以实现复合运算的运算数的复杂转换，以推导融合运算的运算数。请参阅 Keras LSTM。以转换代码为例。

转换为 TensorFlow Lite

使用 TFLiteConverter.from_saved_model API 转换为 TensorFlow Lite。

深入了解

现在，我们将介绍在 TensorFlow Lite 中转换为融合操作的总体设计详情。

TensorFlow 中的编写操作

结合使用 tf.function 和 experimental_implements 函数属性可让用户使用 TensorFlow 原始操作明确编写新操作，并指定生成的复合操作实现的接口。这非常有用，因为它可以提供：

  1. 为底层 TensorFlow 图中的复合操作明确定义的边界。
  2. 明确指定此操作实现的接口。tf.function 的参数与此接口的参数相对应。

例如，我们考虑一个为实现嵌入查找而定义的复合操作。这对应于 TensorFlow Lite 中的融合操作。

  @tf.function(
        experimental_implements="embedding_lookup")
    def EmbFprop(embs, ids_vec):
      """Embedding forward prop.

      Effectively, it computes:
        num = size of ids_vec
        rets = zeros([num, embedding dim])
        for i in range(num):
          rets[i, :] = embs[ids_vec[i], :]
        return rets

      Args:
        embs: The embedding matrix.
        ids_vec: A vector of int32 embedding ids.

      Returns:
        The result of embedding lookups. A matrix of shape
        [num ids in ids_vec, embedding dims].
      """
      num = tf.shape(ids_vec)[0]
      rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))

      def EmbFpropLoop(i, embs, ids_vec, rets):
        # row_id = ids_vec[i]
        row_id = tf.gather(ids_vec, i)
        # row = embs[row_id]
        row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
        # rets[i] = row
        rets = inplace_ops.alias_inplace_update(rets, [i], row)
        return embs, ids_vec, rets

      _, _, rets = functional_ops.For(
          start=0,
          limit=num,
          delta=1,
          inputs=[embs, ids_vec, rets],
          body=EmbFpropLoop,
          rewrite_with_while=compiled)
      if len(weight_shape) > 2:
        rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
      return rets

通过使模型通过 tf.function 使用复合操作（如上所示），可以构建通用基础架构来识别此类操作并将其转换为融合 TensorFlow Lite 操作。

扩展 TensorFlow Lite 转换器

今年早些时候发布的 TensorFlow Lite 转换器仅支持以图表形式导入 TensorFlow 模型，并将所有变量替换为相应的常量值。这不适用于运算融合，因为此类图内联了所有函数，以便变量可以转换为常量。

为了在转换过程中将 tf.function 与 experimental_implements 功能结合使用，需要将这些函数一直保留到转换过程的后期。

因此，我们在转换器中实现了导入和转换 TensorFlow 模型的新工作流，以支持复合运算融合用例。具体而言，添加的新功能包括：

  1. 将 TensorFlow 已保存的模型导入 MLIR
  2. 融合复合操作
  3. 变量可变性分析
  4. 冻结所有只读变量

这样，我们就可以在函数内联和变量冻结之前使用表示复合操作的函数来执行运算融合。

实现运算融合

我们来详细了解一下操作融合传递。此卡券会执行以下操作：

  1. 循环遍历 MLIR 模块中的所有函数。
  2. 如果函数具有 tf._implements 属性，则根据属性值调用相应的操作融合实用程序。
  3. 运算融合实用程序对函数的运算数和属性（充当转换的接口）执行操作，并将函数的正文替换为包含融合运算的等效函数正文。
  4. 在许多情况下，被替换的正文会包含融合操作以外的操作。这些操作对应于函数运算数的一些静态转换，以获取融合运算的运算数。由于这些计算都可以常量折叠起来，因此它们不会出现在只存在融合操作的导出的 FlatBuffer 中。

以下是卡券中的代码段，其中显示了主工作流：

void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
                                                        StringAttr attr) {
  if (attr.getValue() == "embedding_lookup") {
    func.eraseBody();
    func.addEntryBlock();
    // Convert the composite embedding_lookup function body to a
    // TFLite fused embedding_lookup op.
    ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
    if (failed(convert_embedded_lookup.VerifySignature())) {
      return signalPassFailure();
    }
    convert_embedded_lookup.RewriteFunc();
  } else if (attr.getValue() == mlir::TFL::kKerasLstm) {
     func.eraseBody();
     func.addEntryBlock();
     OpBuilder builder(func.getBody());
     if (failed(ConvertKerasLSTMLayer(func, &builder))) {
       return signalPassFailure();
     }
  } else if (.....) /* Other fusions can plug in here */
}

以下代码段显示了将此复合操作映射到 TensorFlow Lite 中将该函数用作转换接口的融合操作。

void RewriteFunc() {
    Value lookup = func_.getArgument(1);
    Value value = func_.getArgument(0);
    auto output_type = func_.getType().getResult(0);

    OpBuilder builder(func_.getBody());
    auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
        func_.getLoc(), output_type, lookup, value);

    builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
  }