TensorFlow 操作融合

概述

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

此基础设施的一个示例用法是 TensorFlow RNN 操作融合到 TensorFlow Lite，如 此处 所述。

什么是融合操作

drawing

TensorFlow 操作可以是原始操作，例如 tf.add，也可以由其他原始操作组成，例如 tf.einsum。原始操作在 TensorFlow 图中显示为单个节点，而复合操作是 TensorFlow 图中的一组节点。执行复合操作等效于执行其所有组成原始操作。

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

融合操作的优势

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

融合操作还提供了一个更高级别的接口来定义复杂的转换，例如量化，否则在更细粒度的级别上将不可行或非常难以实现。

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

融合操作的挑战

将 TensorFlow 中的复合操作转换为 TensorFlow Lite 中的融合操作是一个难题。这是因为

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

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

因此，融合操作的转换已被证明非常具有挑战性。

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

在许多情况下，模型的某些部分可以映射到 TFLite 中的单个操作。这有助于在为特定操作编写优化实现时提高性能。为了能够在 TFLite 中创建融合操作，请识别表示融合操作的图的一部分，并将其包装在 tf.function 中，该函数具有 "experimental_implements" 属性，指向 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 自定义操作 - 请参阅 说明。

请注意，用于注册操作的名称应与 "implements" 签名中指定的 "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 融合操作的总体架构如下

drawing

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

在 TensorFlow 模型源代码中，识别并抽象出复合操作，将其转换为带有 experimental_implements 函数注释的 tf.function。请参阅 嵌入查找 的示例。该函数定义了接口，其参数应用于实现转换逻辑。

编写转换代码

转换代码是根据带有 implements 注释的函数的接口编写的。请参阅 嵌入查找 的融合示例。从概念上讲，转换代码用融合的实现替换了该接口的复合实现。

在 prepare-composite-functions 传递中，插入您的 转换代码。

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

转换为 TensorFlow Lite

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

幕后

我们现在将描述将 TensorFlow 复合操作转换为 TensorFlow Lite 融合操作的总体设计的高级细节。

在 TensorFlow 中组合操作

使用带有 experimental_implements 函数属性的 tf.function 允许用户使用 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 模型导入为图形，其中所有变量都替换为其相应的常量值。这对于操作融合不适用，因为此类图形已将所有函数内联，以便可以将变量转换为常量。

为了在转换过程中利用带有 experimental_implements 特性的 tf.function，需要在转换过程的后期保留这些函数。

因此，我们在转换器中实现了导入和转换 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());
  }