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

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

Toán tử kết hợp là gì

vẽ

Các thao tác của TensorFlow có thể là các hoạt động nguyên gốc (ví dụ: tf.add) hoặc có thể được kết hợp từ các thao tác nguyên gốc khác, chẳng hạn như tf.einsum. Toán tử gốc hiển thị dưới dạng một nút duy nhất trong biểu đồ TensorFlow, còn toán tử 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 thao tác tổng hợp tương đương với việc thực thi từng thao tác chính cấu thành.

Toán tử kết hợp tương ứng với một toán tử duy nhất chịu toàn bộ phép tính do mỗi toán tử chính thực hiện trong toán tử kết hợp tương ứng.

Lợi ích của các phép toán kết hợp

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

Các thao tác kết hợp 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, chẳng hạn như lượng tử hoá, mà nếu ở cấp độ chi tiết hơn thì sẽ không khả thi hoặc rất khó thực hiện.

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

Thách thức với thao tác kết hợp

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

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

  2. Có thể có nhiều phương thức triển khai TensorFlow nhắm đến một thao tác kết hợp của TensorFlow Lite. Ví dụ: có nhiều phương thức triển khai LSTM trong TensorFlow (Keras, CameraXfish/lingvo, v.v.) và mỗi phương thức triển khai này bao gồm các phép toán nguyên thuỷ khác nhau, nhưng tất cả vẫn có thể được chuyển đổi thành cùng một thao tác LSTM kết hợp trong TensorFlow Lite.

Do đó, việc chuyển đổi các toán tử kết hợp đã cho thấy khá khó khăn.

Gói toán tử kết hợp trong tf.function

Trong nhiều trường hợp, một số phần của mô hình có thể được liên kết tới một thao tác duy nhất trong TFLite. Điều này có thể giúp tăng hiệu suất khi viết phương thức triển khai được tối ưu hoá cho các hoạt động cụ thể. Để có thể tạo một thao tác kết hợp trong TFLite, hãy xác định phần biểu đồ đại diện cho thao tác kết hợp rồi gói toán tử đó trong một tf.function với thuộc tính "experimental_implementations" thuộc một tf.function, có giá trị thuộc tính tfl_fusable_op với giá trị true. Nếu thao tác tuỳ chỉnh lấy các thuộc tính thì hãy truyền các thuộc tính đó dưới dạng một "experimental_implementations" tương tự.

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

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

Triển khai hoạt động tuỳ chỉnh và đăng ký bằng Phiên dịch viên TFLite

Triển khai thao tác kết hợp dưới dạng thao tác TFLite Tuỳ chỉnh – xem instructions.

Xin lưu ý rằng tên để đăng ký hoạt động phải tương tự với tên đã chỉ định trong thuộc tính name trong chữ ký triển khai.

Ví dụ về op trong ví dụ này 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ừ phép toán tổng hợp sang phép toán kết hợp (Nâng cao)

Dưới đây là cấu trúc tổng thể để chuyển đổi các thao tác kết hợp của TensorFlow thành các toán tử kết hợp của TensorFlow Lite:

vẽ

Gói toán tử kết hợp trong tf.function

Trong mã nguồn của mô hình TensorFlow, hãy xác định và rút gọn thao tác kết hợp thành tf.function với chú thích hàm experimental_implements. Xem ví dụ về nội dung tra cứu được nhúng. Hàm xác định giao diện và các đối số của giao diện đó sẽ được 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 có chú thích implements. Xem ví dụ về kết hợp tra cứu bằng cách nhúng. Về mặt lý thuyết, mã chuyển đổi sẽ thay thế phương thức triển khai kết hợp của giao diện này bằng giao diện kết hợp.

Trong lượt truyền hàm chuẩn bị kết hợp, hãy bổ trợ vào mã chuyển đổi.

Trong các trường hợp sử dụng nâng cao hơn, bạn có thể triển khai các phép biến đổi phức tạp cho toán hạng của toán tử kết hợp để lấy toán hạng của phép toán kết hợp. Hãy xem ví dụ về mã chuyển đổi Keras LSTM.

Chuyển đổi sang TensorFlow Lite

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

Tìm hiểu sâu

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

Kết hợp thao tác trong TensorFlow

Việc sử dụng tf.function với thuộc tính hàm experimental_implements cho phép người dùng kết hợp rõ ràng các thao tác mới bằng các thao tác gốc của TensorFlow và chỉ định giao diện mà thao tác tổng hợp kết quả sẽ 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 thao tác kết hợp trong biểu đồ TensorFlow cơ bản.
  2. Chỉ định rõ ràng giao diện mà thao tác này triển khai. 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 cùng xem xét một thao tác tổng hợp được xác định để triển khai tính năng tra cứu nhúng. Điều này liên kết với thao tác kết hợp 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 phép toán kết hợp thông qua tf.function như minh hoạ ở trên, bạn có thể xây dựng một cơ sở hạ tầng chung để xác định và chuyển đổi các thao tác đó thành các thao tác kết hợp của TensorFlow Lite.

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

Trình chuyển đổi TensorFlow Lite đã ra mắt vào đầ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 giá trị hằng số tương ứng. Cách này không hiệu quả trong việc hợp nhất phép toán vì các biểu đồ như vậy có tất cả các hàm cùng dòng để 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 một quy trình mới để nhập và chuyển đổi các mô hình TensorFlow trong trình chuyển đổi để hỗ trợ trường hợp sử dụng hợp nhất hoạt động tổng hợp. Cụ thể, các tính năng mới được thêm là:

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

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

Triển khai quá trình hợp nhất thao tác

Hãy cùng tìm hiểu chi tiết hơn về quá trình chuyển đổi thao tác. Thẻ và vé này có những chức năng 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._implementations, dựa trên giá trị thuộc tính, sẽ gọi tiện ích hợp nhất thao tác thích hợp.
  3. Tiện ích hợp nhất toán tử 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 chuyển đổi) và thay thế phần nội dung của hàm bằng một nội dung hàm tương đương chứa toán tử kết hợp.
  4. Trong nhiều trường hợp, phần nội dung được thay thế sẽ chứa các thao tác khác với thao tác kết hợp. Các phép biến đổi này tương ứng với một số phép biến đổi tĩnh trên các toán hạng của hàm để có được toán hạng của phép toán kết hợp. Vì tất cả các phép tính này đều có thể thu gọn không đổi, nên chúng sẽ không xuất hiện trong vùng đệm phẳng được xuất, nơi chỉ có phép toán kết hợp tồn tại.

Đây là đoạn mã của thẻ và vé cho thấy 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 */
}

Dưới đây là đoạn mã cho thấy việc ánh xạ toán tử kết hợp này với một toán tử kết hợp trong TensorFlow Lite sử dụng hàm 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());
  }