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

Ví dụ về cách sử dụng cơ sở hạ tầng này là hợp nhất thao tác TensorFlow RNN với LiteRT, như được trình bày chi tiết tại đây.

Thế nào là các thao tác kết hợp

vẽ

Các thao tác TensorFlow có thể là các thao tác nguyên thuỷ, ví dụ: tf.add hoặc có thể được tạo thành từ các thao tác nguyên thuỷ khác, ví dụ: tf.einsum. Một thao tác nguyên thuỷ xuất hiện dưới dạng một nút duy nhất trong biểu đồ TensorFlow, trong khi một thao tác kết hợp là một tập hợp các nút trong biểu đồ TensorFlow. Thực thi một thao tác kết hợp tương đương với việc thực thi từng thao tác nguyên thuỷ cấu thành.

Một thao tác kết hợp tương ứng với một thao tác duy nhất bao gồm tất cả các phép tính do mỗi thao tác nguyên thuỷ thực hiện trong thao tác kết hợp tương ứng.

Lợi ích của các thao tác kết hợp

Các thao tác kết hợp tồn tại để tối đa hoá hiệu suất của các triển khai nhân cơ bản, bằng cách tối ưu hoá quá trình tính toán tổng thể và giảm mức sử dụng bộ nhớ. Điều này rất có giá trị, đặc biệt đối với các khối lượng công việc suy luận có độ trễ thấp và các 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 như lượng tử hoá. Nếu không, bạn sẽ không thể thực hiện hoặc rất khó thực hiện ở cấp độ chi tiết hơn.

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

Thách thức liên quan đến các hoạt động kết hợp

Việc chuyển đổi các thao tác kết hợp từ TensorFlow sang các thao tác kết hợp trong LiteRT là một vấn đề khó khăn. Nguyên nhân là do:

  1. Các thao tác kết hợp được biểu thị trong biểu đồ TensorFlow dưới dạng một tập hợp các thao tác nguyên thuỷ 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) đồ thị con tương ứng với một thao tác kết hợp như vậy có thể rất khó khăn.

  2. Có thể có nhiều quy trình triển khai TensorFlow nhắm đến một thao tác LiteRT kết hợp. 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 thao tác 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 LiteRT.

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

Gói thao tác kết hợp trong một tf.function

Trong nhiều trường hợp, một phần của mô hình có thể được liên kết vớ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 quy trình 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 của biểu đồ đại diện cho một thao tác kết hợp và gói phần đó trong một tf.function có thuộc tính "experimental_implements" thành 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, hãy truyền chúng như một phần của cùng một "experimental_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.

Triển khai thao tác tuỳ chỉnh và đăng ký bằng Trình thông dịch TFLite

Triển khai thao tác kết hợp dưới dạng một Thao tác tuỳ chỉnh TFLite – xem hướng dẫn.

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

Ví dụ về thao tác 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ừ thao tác kết hợp sang thao tác hợp nhất (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 TensorFlow thành các thao tác kết hợp LiteRT:

vẽ

Gói thao tác kết 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 hoá thao tác kết hợp thành một tf.function bằng chú thích hàm experimental_implements. Xem ví dụ về tra cứu thông tin nhúng. Hàm này xác định giao diện và các đối số của giao diện này phải được dùng để triển khai logic chuyển đổi.

Viết mã chuyển đổi

Mã chuyển đổi được viết theo giao diện của hàm có chú giải implements. Xem ví dụ về việc kết hợp để nhúng tính năng tra cứu. Về mặt lý thuyết, mã chuyển đổi sẽ thay thế việc triển khai thành phần của giao diện này bằng mã chuyển đổi hợp nhất.

Trong giai đoạn prepare-composite-functions, hãy cắm mã chuyển đổi của bạn.

Trong các cách 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 của toán hạng trong thao tác kết hợp để lấy toán hạng của thao tác kết hợp. Hãy xem Keras LSTM. Mã chuyển đổi làm ví dụ.

Chuyển đổi sang LiteRT

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

Tìm hiểu sâu

Giờ đây, chúng ta sẽ mô tả chi tiết cấp cao về thiết kế tổng thể trong việc chuyển đổi sang các hoạt động kết hợp trong LiteRT.

Kết hợp các 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ách sử dụng các thao tác nguyên thuỷ của TensorFlow và chỉ định giao diện mà thao tác kết hợp thu được sẽ triển khai. Đ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õ cho hoạt động kết hợp trong biểu đồ TensorFlow cơ bản.
  2. Chỉ định rõ 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 xem xét một thao tác kết hợp được xác định để triển khai tính năng tra cứu thông tin nhúng. Thao tác này tương ứng với một thao tác kết hợp trong 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

Bằng cách tạo các mô hình sử dụng các thao tác 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 như vậy thành các thao tác LiteRT kết hợp.

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

Trình chuyển đổi LiteRT đượ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 một biểu đồ với tất cả các biến được thay thế bằng các giá trị hằng số tương ứng. Điều này không hoạt động đối với việc hợp nhất các thao tác vì những đồ 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 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ữ lại 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 nhằm hỗ trợ trường hợp sử dụng hợp nhất thao tác kết 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 của TensorFlow vào MLIR
  2. fuse composite operations
  3. phân tích khả năng thay đổi của biến
  4. đóng băng tất cả các biến chỉ có thể đọc

Điều này cho phép chúng ta 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 thao tác kết hợp trước khi nội tuyến hoá hàm và cố định biến.

Triển khai tính năng hợp nhất thao tác

Hãy cùng tìm hiểu kỹ hơn về thẻ kết hợp thao tác. Thẻ này có những chức năng sau:

  1. Lặp lại tất cả các hàm 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, hãy gọi tiện ích hợp nhất thao tác thích hợp.
  3. Tiện ích hợp nhất thao tác 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 quá trình chuyển đổi) và thay thế nội dung của hàm bằng một nội dung 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 nội dung được thay thế sẽ chứa các thao tác khác ngoài thao tác kết hợp. Những thao tác 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 để thu được các toán hạng của thao tác hợp nhất. Vì tất cả các phép tính này đều có thể được giảm bớt hằng số, nên chúng sẽ không xuất hiện trong flatbuffer đã xuất, nơi chỉ có hoạt động kết hợp tồn tại.

Dưới đây là đoạn mã từ thẻ và vé cho thấy quy trình 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 cách ánh xạ thao tác kết hợp này thành một thao tác kết hợp trong LiteRT bằng cách tận 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());
  }