TensorFlow 運算融合

總覽

本頁面說明將 TensorFlow 複合運算轉換為 TensorFlow Lite 中融合運算所需的設計和步驟。這個基礎架構一般具有用途，支援將 TensorFlow 中的任何複合運算轉換為 TensorFlow Lite 中的對應融合運算。

這個基礎架構使用的範例為 TensorFlow RNN 作業融合至 TensorFlow Lite，詳情請參閱這裡。

什麼是融合運算

繪畫

TensorFlow 作業可以是原始運算 (例如 tf.add)，也可以透過其他原始作業 (例如 tf.einsum) 組成。基本作業在 TensorFlow 圖表中顯示為單一節點，而複合作業則是 TensorFlow 圖形中的一組節點。執行複合運算相當於執行其每個組成基本作業。

融合運算與單一作業相對應，該運算會耗用對應複合式作業中每個原始作業執行的所有運算。

融合運算的優點

融合作業可最佳化整體運算作業並減少記憶體使用量，藉此改善基礎核心實作的效能。這非常寶貴，對於低延遲的推論工作負載和資源受限的行動平台而言更是如此。

融合運算也提供更高級別的介面，用於定義量化等複雜轉換作業，否則這些是無法實現，或難以更精細的執行。

基於上述原因，TensorFlow Lite 有許多融合運算的執行個體。這些融合作業通常會對應來源 TensorFlow 程式中的複合作業。TensorFlow 中以單一融合運算的形式在 TensorFlow 中實作的複合運算範例包括各種 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_implementations」屬性至 tf.function，且其屬性值 tfl_fusable_op 為 true。如果自訂作業採用屬性，請將其做為同一個「experimental_implementations」的一部分傳遞。

例如：

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 translateer 註冊

以 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 註解的函式介面編寫。請參閱嵌入查詢的融合範例。從概念上來說，轉換程式碼會以整合式程式碼取代這個介面的複合實作。

在準備複合函式傳遞中，將轉換程式碼中的外掛程式加入外掛程式。

若是更進階的用途，可以實作複合式運算運算元的複雜轉換，以衍生融合運算的運算元。請參閱 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 模型匯入為圖表，且所有變數都會替換為對應的常數值。這不適用於作業融合，因為此類圖表會內嵌所有函式，以便變數可以轉換成常數。

為了在轉換程序期間搭配 experimental_implements 功能使用 tf.function，需要保留函式到後續的轉換程序。

因此，我們導入了新的工作流程，以便在轉換工具中匯入及轉換 TensorFlow 模型，以支援複合運算融合用途。具體來說，新增的功能包括：

  1. 將 TensorFlow 儲存的模型匯入 MLIR
  2. 融合複合運算
  3. 可變動性分析
  4. 凍結所有唯讀變數

這可讓我們先使用代表複合運算的函式執行運算融合，再進行函式內嵌和變數凍結。

實作作業融合

一起來深入瞭解作業融合傳遞。此票證會執行以下作業：

  1. 循環瀏覽 MLIR 模組中的所有函式。
  2. 如果函式根據屬性值提供 tf._implementations 屬性，請呼叫適當的運算 fusion 公用程式。
  3. 作業融合公用程式會對函式的運算元和屬性 (做為轉換的介面) 運作，並將函式主體替換為包含融合運算的對等函式主體。
  4. 在許多情況下，替換的主體會包含融合運算以外的作業。這些參數會對應至函式運算元的部分靜態轉換，以取得融合運算的運算元。由於這些運算都可以一律折疊，因此不會顯示在匯出的扁平緩衝區中，而只有融合運算。

以下是票證中的程式碼片段，其中顯示主要工作流程：

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