隆重推出 LiteRT：Google 為裝置端 AI (舊稱 TensorFlow Lite) 打造的高效能執行階段。

本頁面由 Cloud Translation API 翻譯而成。

TensorFlow 運算融合

總覽

本頁說明轉換複合作業所需的設計和步驟以及 LiteRT 中的融合作業這種基礎架構一般用途，並支援在 TensorFlow 中轉換任何複合運算與 LiteRT 中的對應融合運算做比較

TensorFlow RNN 運算融合至這個基礎架構 LiteRT，詳情請參閱這篇文章。

什麼是融合運算？

TensorFlow 運算可以是原始作業，例如： tf.add 也可以源自其他原始運算，例如 tf.einsum。某個原始作業會在 TensorFlow 圖形中顯示為單一節點，「Operation」是 TensorFlow 圖形中的一組節點執行複合運算相當於執行每個組成基元作業。

融合運算會對應於置入所有每個原始運算所執行的複合運算。

融合作業的優點

有融合作業是為了最大化基礎核心的效能方法是最佳化整體運算作業及減少記憶體調整規模。這非常寶貴，尤其是對於低延遲推論工作負載和資源受限的行動平台

融合作業也提供更高層級的介面，用於定義複雜作業像量化這類模型，效果不彰難以執行更精細的任務

基於原因，LRTRT 有許多融合作業例項。這些融合運算通常對應於複合原始碼中的運算作業中的複合運算範例在 LiteRT 中以單一融合運算的形式實作的 TensorFlow 包括各種 RNN 運算，例如單向和雙向序列 LSTM、卷積 (conv2d、偏誤、relu)、完全連結 (matmul、偏誤增加、 relu) 和其他頻道。LiteRT 中的 LSTM 量化功能目前僅適用於融合 LSTM 行動中

融合作業的挑戰

將複合運算從 TensorFlow 轉換為融合運算 LiteRT 是一大難題。可能的原因如下：

在 TensorFlow 圖形中，複合運算是以一組未定義界線的原始作業。可以難以辨識 (例如透過模式比對) 針對這種複合運算進行對應。
針對 Fused 為目標的 TensorFlow 實作項目可能不只一個 LiteRT 作業。例如許多 LSTM 實作 TensorFlow (Keras、Babelfish/lingvo 等) 都由模型組成不同的原始運算，但全部仍可轉換為相同的融合 LSTM 運算

因此，已證實融合融合營運的困難性。

從複合運算轉換為 TFLite 自訂作業 (建議)

將複合運算納入 `tf.function`

在許多情況下，模型的某些部分可以對應至 TFLite.這有助於編寫最佳化的實作方式來提高成效並用於特定作業要在 TFLite 中建立融合作業找出圖表中整合融合運算的各部分 tf.function 「實驗實作」設為 tf.function，且該屬性含有值為 true，值為 tfl_fusable_op。如果自訂作業屬性，然後將其做為相同「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 解譯器註冊

將融合作業實作為 TFLite 自訂作業 - 請參閱操作說明。

請注意，註冊運算的名稱應與名稱相似在實作簽章的 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 複合運算轉換為 LiteRT 融合運算如下：

將複合運算納入 `tf.function`

在 TensorFlow 模型原始碼中，識別並提取複合資料轉換為 tf.function experimental_implements 函式註解。請參閱嵌入查詢範例。函式會定義介面，並使用其引數來實作轉換邏輯

撰寫轉換程式碼

轉換程式碼是根據使用 implements 註解。查看嵌入功能的範例 lookup。概念上，轉換程式碼會取代複合程式碼實作這個介面與整合式介面。

在 prepare-複合函式傳遞中，將轉換中的外掛程式程式碼。

在更進階的用法中，可以看到複合運算的運算元，用來衍生融合的運算元作業。查看 Keras LSTM. 轉換程式碼做為範例

轉換為 LiteRT

使用 TFLiteConverter.from_saved_model 要轉換成 LiteRT 的 API。

深入解析

現在我們針對從模型轉換為融合設計的整體設計，提供概略的細節以及 LiteRT 中的作業

在 TensorFlow 中編寫作業

使用 tf.function 使用 experimental_implements 函式屬性可讓使用者使用函式 TensorFlow 原始作業，並指定產生結果的介面複合運算實作這項功能非常有用，因為：

基礎模型中，複合式作業的定義明確界線 TensorFlow 圖表
請明確指定這項作業實作的介面。引數 tf.function 會對應這個介面的引數

舉例來說，假設定義為實作的複合運算嵌入查詢這會對應至 LiteRT 中的融合作業。

  @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 就能建構一般基礎架構 識別並轉換這類作業，並轉換為整合式 LiteRT 作業。

擴充 LiteRT 轉換工具

僅支援今年稍早推出的 LiteRT 轉換工具將 TensorFlow 模型匯入為圖表，其中所有變數都會替換為對應的常數值。這就不適用於融合作業，因為這些圖形全都具有內嵌函式，因此您可以將這些變數轉換為常數。

為了使用 tf.function experimental_implements 功能，直到轉換程序的後續階段為止

因此，我們導入新的工作流程，匯入及轉換 TensorFlow 模型中，藉此支援複合運算融合用途具體來說，新增的功能如下：

以便我們使用能代表複合運算

實作運算融合

接著就來詳細瞭解融合作業流程。這麼做會包括：

循環執行 MLIR 模組中的所有函式。
如果函式具有 tf._implementations 屬性，則根據屬性值時，呼叫適當的作業融合公用程式。
運算融合公用程式會在函式的運算元上屬性 (做為轉換介面) 並取代相應函式主體的函式主體融合作業。
在許多情況下，取代主體包含的作業融合作業。這些字串對應至了函式的運算元，用於取得融合運算的運算元。由於這些計算可以持續折疊，因此不會存在於匯出的扁平緩衝區中，只有融合作業個值。

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

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 */
}

下列程式碼片段顯示將這項複合作業對應至融合作業作業，利用函式做為轉換介面。

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

TensorFlow 運算融合

總覽

什麼是融合運算？

融合作業的優點

融合作業的挑戰

從複合運算轉換為 TFLite 自訂作業 (建議)

將複合運算納入 tf.function

實作自訂運算，並使用 TFLite 解譯器註冊

從複合作業轉換為融合運算 (進階)

將複合運算納入 tf.function

撰寫轉換程式碼

轉換為 LiteRT

深入解析

在 TensorFlow 中編寫作業

擴充 LiteRT 轉換工具

實作運算融合

將複合運算納入 `tf.function`

將複合運算納入 `tf.function`