Kết hợp hoạt động TensorFlow, kết hợp hoạt động TensorFlow

Sử dụng bộ sưu tập để sắp xếp ngăn nắp các trang Lưu và phân loại nội dung dựa trên lựa chọn ưu tiên của bạn.

Tổng quan

Trang này mô tả thiết kế và các bước cần thiết để chuyển đổi các hoạt động tổng hợp trong TensorFlow thành các hoạt động hợp nhất trong TensorFlow Lite. Cơ sở hạ tầng này là mục đích chung và hỗ trợ chuyển đổi bất kỳ hoạt động tổng hợp nào trong TensorFlow sang hoạt động hợp nhất tương ứng trong TensorFlow Lite.

Một ví dụ sử dụng cơ sở hạ tầng này là kết hợp hoạt động TensorFlow RNN thành TensorFlow Lite, như chi tiết ở đây .

Hoạt động hợp nhất là gì

đang vẽ

Các hoạt động TensorFlow có thể là các hoạt động nguyên thủy, ví dụ tf.add hoặc chúng có thể được tạo từ các hoạt động nguyên thủy khác, ví dụ như tf.einsum . Một hoạt động nguyên thủy hiển thị dưới dạng một nút duy nhất trong biểu đồ TensorFlow trong khi hoạt động tổng hợp là một tập hợp các nút trong biểu đồ TensorFlow. Việc thực thi một phép toán tổng hợp tương đương với việc thực thi mỗi phép toán nguyên thủy cấu thành của nó.

Một phép toán hợp nhất tương ứng với một phép toán đơn lẻ bao gồm tất cả các phép tính được thực hiện bởi mỗi phép toán nguyên thủy trong phép toán tổng hợp tương ứng.

Lợi ích của hoạt động hợp nhất

Các hoạt động hợp nhất tồn tại để tối đa hóa hiệu suất của các triển khai nhân cơ bản của chúng, bằng cách tối ưu hóa tính toán tổng thể và giảm dung lượng bộ nhớ. Điều này rất có giá trị, đặc biệt đối với khối lượng công việc suy luận có độ trễ thấp và các nền tảng di động hạn chế tài nguyên.

Các phép toán hợp nhất cũng cung cấp một giao diện cấp cao hơn để xác định các phép biến đổi phức tạp như lượng tử hóa, điều này sẽ không khả thi hoặc rất khó thực hiện ở cấp độ chi tiết hơn.

TensorFlow Lite có nhiều trường hợp hoạt động hợp nhất vì những lý do đã nêu ở trên. Các hoạt động hợp nhất này thường tương ứng với các hoạt động tổng hợp trong chương trình TensorFlow nguồn. Ví dụ về các hoạt động tổng hợp trong TensorFlow được triển khai dưới dạng một hoạt động hợp nhất duy nhất trong TensorFlow Lite bao gồm các hoạt động RNN khác nhau như LSTM chuỗi đơn hướng và hai chiều, tích chập (chập 2, thêm thiên vị, relu), kết nối đầy đủ (matmul, thêm thiên vị, relu) và hơn thế nữa . Trong TensorFlow Lite, lượng tử hóa LSTM hiện chỉ được thực hiện trong các hoạt động LSTM hợp nhất.

Những thách thức với hoạt động hợp nhất

Chuyển đổi các hoạt động tổng hợp từ TensorFlow sang các hoạt động hợp nhất trong TensorFlow Lite là một vấn đề khó khăn. Điều này là do:

  1. Các phép toán tổng hợp được biểu diễn trong đồ thị TensorFlow dưới dạng một tập hợp các phép toán nguyên thủy không có ranh giới xác định rõ. Có thể rất khó để xác định (ví dụ: thông qua đối sánh mẫu) biểu đồ phụ tương ứng với một phép toán tổng hợp như vậy.

  2. Có thể có nhiều hơn một triển khai TensorFlow nhắm mục tiêu một hoạt động TensorFlow Lite được hợp nhất. Ví dụ: có rất nhiều triển khai LSTM trong TensorFlow (Keras, Babelfish / lingvo, v.v.) và mỗi hoạt động này bao gồm các hoạt động nguyên thủy khác nhau nhưng tất cả chúng vẫn có thể được chuyển đổi thành cùng một hoạt động LSTM được hợp nhất trong TensorFlow Lite.

Do đó, việc chuyển đổi các hoạt động hợp nhất đã được chứng minh là khá khó khăn.

Kết hợp hoạt động tổng hợp trong một tf.function

Trong nhiều trường hợp, một số phần của mô hình có thể được ánh xạ tới một hoạt động duy nhất trong TFLite. Điều này có thể giúp tăng hiệu suất khi viết một triển khai được tối ưu hóa cho các hoạt động cụ thể. Để có thể tạo một hoạt động hợp nhất trong TFLite, hãy xác định phần của biểu đồ biểu thị một hoạt động được hợp nhất và bọc nó trong một tf.function với thuộc tính "Experiment_implements" thành một hàm tf. tfl_fusable_op tf.function giá trị true . Nếu thao tác tùy chỉnh nhận thuộc tính thì hãy chuyển chúng như một phần của cùng một "thử nghiệm_implements".

Thí dụ,

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

Lưu ý rằng bạn không cần đặt allow_custom_ops trên trình chuyển đổi làm thuộc tính tfl_fusable_op ngụ ý điều này đã có.

Triển khai tùy chọn tùy chỉnh và đăng ký với Trình thông dịch TFLite

Triển khai hoạt động hợp nhất của bạn dưới dạng hoạt động Tùy chỉnh TFLite - xem hướng dẫn .

Lưu ý rằng, tên để đăng ký op phải giống với tên được chỉ định trong thuộc tính name trong chữ ký triển khai.

Một ví dụ cho op trong ví dụ là

  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 coder.
    return kTfLiteOk;
  };
  reg.builtin_code = kTfLiteCustom;
  resolver->AddCustom(kOpName, &reg);

Chuyển đổi từ hỗn hợp sang hoạt động hợp nhất (Nâng cao)

Kiến trúc tổng thể để chuyển đổi các hoạt động kết hợp TensorFlow thành các hoạt động kết hợp TensorFlow Lite như sau:

đang vẽ

Kết hợp hoạt động tổng hợp trong một tf.function

Trong mã nguồn của mô hình TensorFlow, hãy xác định và tóm tắt phép toán tổng hợp thành một tf.function với chú thích hàm Experiment_implements . Xem ví dụ về tra cứu nhúng . Hàm xác định giao diện và các đối số của nó nên được sử dụng để triển khai logic chuyển đổi.

Viết mã chuyển đổi

Mã chuyển đổi được viết trên giao diện của hàm với chú thích implements . Xem kết hợp ví dụ để tra cứu nhúng . Về mặt khái niệm, mã chuyển đổi thay thế việc triển khai kết hợp của giao diện này bằng giao diện được hợp nhất.

Trong thẻ chuẩn bị-tổng hợp-các chức năng, hãy bổ sung plugin trong mã chuyển đổi của bạn.

Trong các cách sử dụng nâng cao hơn, có thể thực hiện các phép biến đổi phức tạp của các toán hạng của phép toán tổng hợp để lấy ra các toán hạng của phép toán hợp nhất. Xem Keras LSTM . mã chuyển đổi làm ví dụ.

Chuyển đổi sang TensorFlow Lite

Sử dụng API TFLiteConverter.from_saved_model để chuyển đổi sang TensorFlow Lite.

Dưới mui xe

Bây giờ chúng tôi mô tả chi tiết cấp cao của thiết kế tổng thể trong việc chuyển đổi sang các hoạt động hợp nhất trong TensorFlow Lite.

Soạn các thao tác trong TensorFlow

Việc sử dụng tf.function với thuộc tính hàm Experiment_implements cho phép người dùng soạn các phép toán mới một cách rõ ràng bằng cách sử dụng các phép toán nguyên thủy TensorFlow và chỉ định giao diện mà phép toán tổng hợp kết quả triển khai. Điều này rất hữu ích vì nó cung cấp:

  1. Ranh giới được xác định rõ ràng cho hoạt động tổng hợp trong biểu đồ TensorFlow bên dưới.
  2. Chỉ định rõ ràng giao diện mà thao tác này thực hiện. Các đối số của hàm tf.function tương ứng với các đối số của giao diện này.

Ví dụ, chúng ta hãy xem xét một hoạt động tổng hợp được xác định để thực hiện tra cứu nhúng. Điều này ánh xạ đến một hoạt động hợp nhất trong 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

Bằng cách làm cho các mô hình sử dụng các hoạt động tổng hợp thông qua tf.function như minh họa ở trên, có thể xây dựng một cơ sở hạ tầng chung để xác định và chuyển đổi các hoạt động đó thành các hoạt động TensorFlow Lite được hợp nhất.

Mở rộng công cụ chuyển đổi TensorFlow Lite

Công cụ chuyển đổi TensorFlow Lite được phát hành vào đầu năm nay chỉ hỗ trợ nhập các mô hình TensorFlow dưới dạng biểu đồ với tất cả các biến được thay thế bằng các giá trị không đổi tương ứng của chúng. Điều này không hoạt động đối với sự kết hợp hoạt động vì các đồ thị như vậy có tất cả các hàm được nội tuyến để các biến có thể được chuyển thành hằng số.

Để tận dụng chức năng tf. experimental_implements tính năng tf.function trong quá trình chuyển đổi, các chức năng cần được bảo toàn cho đến sau này trong quá trình chuyển đổi.

Do đó, chúng tôi đã triển khai quy trình nhập và chuyển đổi mô hình TensorFlow mới trong bộ chuyển đổi để hỗ trợ trường hợp sử dụng kết hợp hoạt động tổng hợp. Cụ thể, các tính năng mới được thêm vào là:

  1. Nhập các mô hình đã lưu TensorFlow vào MLIR
  2. cầu chì hoạt động tổng hợp
  3. phân tích khả năng thay đổi
  4. đóng băng tất cả các biến chỉ đọc

Điều này cho phép chúng tôi thực hiện hợp nhất hoạt động bằng cách sử dụng các hàm đại diện cho các hoạt động tổng hợp trước khi nội tuyến hàm và đóng băng biến.

Thực hiện hợp nhất hoạt động

Hãy xem xét chi tiết hơn về hoạt động hợp hạch. Thẻ này thực hiện những điều sau:

  1. Lặp lại tất cả các chức năng trong mô-đun MLIR.
  2. Nếu một hàm có thuộc tính tf._implements, dựa trên giá trị thuộc tính, sẽ gọi tiện ích kết hợp hoạt động thích hợp.
  3. Tiện ích kết hợp hoạt động hoạt động trên các toán hạng và thuộc tính của hàm (đóng vai trò là giao diện cho việc chuyển đổi) và thay thế phần thân của hàm bằng một phần thân của hàm tương đương có chứa hoạt động được hợp nhất.
  4. Trong nhiều trường hợp, phần thân được thay thế sẽ chứa các hoạt động khác với hoạt động hợp nhất. Các biến đổi này tương ứng với một số biến đổi tĩnh trên các toán hạng của hàm để thu được các toán hạng của hoạt động hợp nhất. Vì tất cả các phép tính này đều có thể được gấp lại liên tục, chúng sẽ không hiện diện trong bộ đệm phẳng được xuất, nơi chỉ tồn tại hoạt động hợp nhất.

Đây là đoạn mã từ thẻ hiển thị quy trình làm việc chính:

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

Đây là đoạn mã hiển thị ánh xạ hoạt động tổng hợp này với một hoạt động hợp nhất trong TensorFlow Lite tận dụng chức năng làm giao diện chuyển đổi.

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