Hợp nhất hoạt động TensorFlow

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 thao tác tổng hợp trong TensorFlow thành các thao tác hợp nhất trong TensorFlow Lite. Cơ sở hạ tầng này phục vụ mục đích chung và hỗ trợ chuyển đổi mọi hoạt động tổng hợp trong TensorFlow sang hoạt động hợp nhất tương ứng trong TensorFlow Lite.

Một ví dụ về cách sử dụng cơ sở hạ tầng này là sự hợp nhất hoạt động TensorFlow RNN với TensorFlow Lite, như được nêu chi tiết tại đây .

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

vẽ

Các hoạt động của TensorFlow có thể là các hoạt động nguyên thủy, ví dụ như tf.add hoặc chúng có thể được tạo thành 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 hiện một thao tác tổng hợp tương đương với việc thực hiện từng thao tác 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ẻ gộp tất cả các tính toán đượ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 việc triển khai kernel cơ bản, 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à nền tảng di động bị hạn chế về tài nguyên.

Các hoạt động hợp nhất cũng cung cấp 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 một chiều và hai chiều, tích chập (conv2d, cộng thiên vị, relu), kết nối đầy đủ (matmul, thiên vị cộng, relu), v.v. . Trong TensorFlow Lite, lượng tử hóa LSTM hiện chỉ được triển khai trong các hoạt động LSTM hợp nhất.

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

Việc chuyển đổi các phép toán tổng hợp từ TensorFlow sang các phép toán hợp nhất trong TensorFlow Lite là một bài toán khó. Điều này là do:

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

  2. Có thể có nhiều triển khai TensorFlow nhắm mục tiêu hoạt động TensorFlow Lite hợp nhất. Ví dụ: có nhiều cách triển khai LSTM trong TensorFlow (Keras, Babelfish/lingvo, v.v.) và mỗi cách triển khai 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 hợp nhất trong TensorFlow Lite.

Như vậy, việc chuyển đổi các hoạt động hợp nhất đã tỏ ra khá khó khăn.

Bao bọc 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 thao tác duy nhất trong TFLite. Điều này có thể giúp cải thiện 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 biểu đồ đại diện cho một hoạt động hợp nhất và gói nó trong một tf.function với thuộc tính "experimental_implements" thành một tf.function , có giá trị thuộc tính tfl_fusable_op với giá trị true . Nếu thao tác tùy chỉnh lấy các thuộc tính thì chuyển chúng như một phần của cùng một "thử nghiệm_implements".

Ví 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 bộ chuyển đổi vì thuộc tính tfl_fusable_op đã ngụ ý điều này rồi.

Triển khai op 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ý thực hiện.

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

Chuyển đổi từ hoạt động tổng 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 tổng hợp TensorFlow sang các hoạt động hợp nhất TensorFlow Lite như sau:

vẽ

Bao bọc hoạt động tổng hợp trong một tf.function

Trong mã nguồn mô hình TensorFlow, hãy xác định và trừu tượng hóa hoạt động tổng hợp thành một tf.function với chú thích hàm Experiment_implements . Xem ví dụ về nhúng tra cứu . Hàm xác định giao diện và các đối số của nó sẽ đượ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 một ví dụ tổng hợp để nhúng tra cứu . Về mặt khái niệm, mã chuyển đổi thay thế việc triển khai tổng hợp giao diện này bằng giao diện hợp nhất.

Trong thẻ chuẩn bị hàm tổng hợp, hãy bổ sung vào 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 toán hạng của phép toán tổng hợp để rút ra 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 về thiết kế tổng thể khi chuyển đổi sang hoạt động hợp nhất trong TensorFlow Lite.

Soạn các hoạt động 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 thảo rõ ràng các hoạt động mới bằng cách sử dụng các hoạt động nguyên thủy của TensorFlow và chỉ định giao diện mà hoạt động tổng hợp kết quả thực hiện. Điều này rất hữu ích vì nó cung cấp:

  1. Một ranh giới được xác định rõ ràng cho hoạt động tổng hợp trong biểu đồ TensorFlow cơ bản.
  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 tf.function tương ứng với các đối số của giao diện này.

Ví dụ: hãy xem xét một thao tác tổng hợp được xác định để triển khai tra cứu nhúng. Điều này ánh xạ tới 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 tạo 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 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 hợp nhất của TensorFlow Lite.

Mở rộng bộ chuyển đổi TensorFlow Lite

Bộ chuyển đổi TensorFlow Lite được phát hành đầu năm nay chỉ hỗ trợ nhập 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 hiệu quả đối với phép hợp nhất vì các biểu đồ 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 tf.function với tính năng experimental_implements trong quá trình chuyển đổi, các hàm cần được giữ nguyên cho đến sau này trong quá trình chuyển đổi.

Do đó, chúng tôi đã triển khai quy trình làm việc mới nhằm nhập và chuyển đổi các mô hình TensorFlow 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ể những tính năng mới được bổ sung là:

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

Điều này cho phép chúng ta 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

Chúng ta hãy xem xét quá trình hợp nhất hoạt động chi tiết hơn. Pass này thực hiện như 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 hợp nhất hoạt động thích hợp.
  3. Tiện ích hợp nhất 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 phần thân hàm tương đương có chứa thao tá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 thao tác khác với thao tác hợp nhất. Chúng tương ứng với một số biến đổi tĩnh trên toán hạng của hàm để thu được toán hạng của phép toán hợp nhất. Vì tất cả các tính toán này đều có thể được gấp lại liên tục nên chúng sẽ không xuất hiệ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 sang hoạt động hợp nhất trong TensorFlow Lite tận dụng chức năng này như một 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());
  }