ฟิวชันการดำเนินการ TensorFlow

ภาพรวม

หน้านี้อธิบายการออกแบบและขั้นตอนที่จำเป็นในการแปลงการดำเนินการแบบคอมโพสิต ใน TensorFlow เป็นการดำเนินการที่ผสานใน LiteRT โครงสร้างพื้นฐานนี้มีวัตถุประสงค์ทั่วไปและรองรับการแปลงการดำเนินการแบบคอมโพสิตใดๆ ใน TensorFlow เป็นการดำเนินการที่ผสานที่สอดคล้องกันใน LiteRT

ตัวอย่างการใช้โครงสร้างพื้นฐานนี้คือการผสานการทำงานของ TensorFlow RNN เข้ากับ LiteRT ดังที่ระบุไว้ที่นี่

การดำเนินการแบบผสานคืออะไร

การวาดภาพ

การดำเนินการ TensorFlow อาจเป็นการดำเนินการดั้งเดิม เช่น tf.add หรืออาจ ประกอบจากการดำเนินการดั้งเดิมอื่นๆ เช่น tf.einsum การดำเนินการแบบดั้งเดิม จะปรากฏเป็นโหนดเดียวในกราฟ TensorFlow ขณะที่การดำเนินการแบบคอมโพสิต คือชุดของโหนดในกราฟ TensorFlow การดำเนินการ แบบคอมโพสิตเทียบเท่ากับการดำเนินการแต่ละรายการที่เป็นองค์ประกอบพื้นฐาน

การดำเนินการที่ผสานจะสอดคล้องกับการดำเนินการเดียวที่ครอบคลุมการคำนวณทั้งหมดที่ดำเนินการโดยการดำเนินการพื้นฐานแต่ละรายการในการดำเนินการแบบคอมโพสิตที่เกี่ยวข้อง

ประโยชน์ของการดำเนินการที่ผสานรวม

การดำเนินการที่ผสานรวมมีไว้เพื่อเพิ่มประสิทธิภาพของการใช้งานเคอร์เนลพื้นฐานให้ได้สูงสุด โดยการเพิ่มประสิทธิภาพการคำนวณโดยรวมและลดการใช้หน่วยความจำ ซึ่งมีประโยชน์อย่างยิ่ง โดยเฉพาะอย่างยิ่งสำหรับภาระงานการอนุมานที่มีเวลาในการตอบสนองต่ำ และแพลตฟอร์มอุปกรณ์เคลื่อนที่ที่มีทรัพยากรจำกัด

นอกจากนี้ การดำเนินการที่ผสานรวมยังมีอินเทอร์เฟซระดับสูงกว่าเพื่อกำหนดการแปลงที่ซับซ้อน เช่น การหาปริมาณ ซึ่งไม่เช่นนั้นอาจทำไม่ได้หรือทำได้ยากมากในระดับที่ละเอียดยิ่งขึ้น

LiteRT มีอินสแตนซ์ของการดำเนินการที่ผสานรวมหลายรายการด้วยเหตุผลที่ระบุไว้ข้างต้น โดยปกติแล้วการดำเนินการที่ผสานรวมเหล่านี้จะสอดคล้องกับการดำเนินการแบบคอมโพสิตในโปรแกรม TensorFlow ต้นทาง ตัวอย่างของการดำเนินการแบบคอมโพสิตใน TensorFlow ที่ มีการใช้งานเป็นการดำเนินการแบบผสานเดียวใน LiteRT ได้แก่ การดำเนินการ RNN ต่างๆ เช่น LSTM แบบลำดับทิศทางเดียวและแบบสองทิศทาง คอนโวลูชัน (conv2d, bias add, relu), การเชื่อมต่อแบบสมบูรณ์ (matmul, bias add, relu) และอื่นๆ ใน LiteRT ขณะนี้การวัดปริมาณ LSTM จะใช้ได้เฉพาะใน LSTM ที่ผสานรวม เท่านั้น

ความท้าทายของการดำเนินการที่ผสานรวม

การแปลงการดำเนินการแบบคอมโพสิตจาก TensorFlow เป็นการดำเนินการแบบผสานใน LiteRT เป็น ปัญหาที่ซับซ้อน ซึ่งเกิดจากสาเหตุต่อไปนี้

  1. การดำเนินการแบบคอมโพสิตจะแสดงในกราฟ TensorFlow เป็นชุดของการดำเนินการดั้งเดิมที่ไม่มีขอบเขตที่กำหนดไว้อย่างชัดเจน การระบุ (เช่น ผ่านการจับคู่รูปแบบ) กราฟย่อยที่สอดคล้องกับการดำเนินการแบบคอมโพสิตดังกล่าวอาจเป็นเรื่องที่ท้าทายมาก

  2. การใช้งาน TensorFlow อาจมีมากกว่า 1 รายการที่กำหนดเป้าหมายเป็นการดำเนินการที่ผสานรวมของ LiteRT ตัวอย่างเช่น มีการใช้งาน LSTM หลายอย่างใน TensorFlow (Keras, Babelfish/lingvo ฯลฯ) และแต่ละอย่างประกอบด้วยการดำเนินการพื้นฐานที่แตกต่างกัน แต่ทั้งหมดนี้ยังคงแปลงเป็นการดำเนินการ LSTM ที่ผสานรวมเดียวกันใน LiteRT ได้

ด้วยเหตุนี้ การแปลงการดำเนินการที่ผสานรวมจึงเป็นสิ่งที่ท้าทายอย่างยิ่ง

รวมการดำเนินการแบบคอมโพสิตไว้ใน tf.function

ในหลายกรณี ส่วนหนึ่งของโมเดลสามารถเชื่อมโยงกับการดำเนินการเดียวใน TFLite ได้ ซึ่งจะช่วยเพิ่มประสิทธิภาพเมื่อเขียนการติดตั้งใช้งานที่เพิ่มประสิทธิภาพ สำหรับการดำเนินการที่เฉพาะเจาะจง หากต้องการสร้างการดำเนินการที่ผสานใน TFLite ให้ระบุส่วนของกราฟที่แสดงการดำเนินการที่ผสานและห่อไว้ใน tf.function ที่มี แอตทริบิวต์ "experimental_implements" เป็น tf.function ซึ่งมีแอตทริบิวต์ ค่า tfl_fusable_op ที่มีค่า true หากการดำเนินการที่กำหนดเองต้องใช้แอตทริบิวต์ ให้ส่งแอตทริบิวต์เหล่านั้นเป็นส่วนหนึ่งของ "experimental_implements" เดียวกัน

ตัวอย่าง

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 หมายถึงการตั้งค่านี้อยู่แล้ว

ใช้ Op ที่กำหนดเองและลงทะเบียนกับ TFLite Interpreter

ใช้การดำเนินการที่ผสานรวมเป็น Custom Operation ของ 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 ดูตัวอย่างการฝังการค้นหา ฟังก์ชัน จะกำหนดอินเทอร์เฟซและควรใช้อาร์กิวเมนต์เพื่อใช้ตรรกะ Conversion

เขียนโค้ด Conversion

โค้ด Conversion เขียนตามอินเทอร์เฟซของฟังก์ชันที่มีคำอธิบายประกอบ implements ดูตัวอย่างการผสานสำหรับการค้นหาการฝัง ในเชิงแนวคิด โค้ด Conversion จะแทนที่การใช้งานแบบคอมโพสิต ของอินเทอร์เฟซนี้ด้วยการใช้งานแบบผสาน

ในพาส prepare-composite-functions ให้เสียบโค้ด Conversion

ในการใช้งานขั้นสูง คุณสามารถใช้การแปลงที่ซับซ้อนของ ตัวถูกดำเนินการของการดำเนินการแบบคอมโพสิตเพื่อหาตัวถูกดำเนินการของการดำเนินการที่ผสานรวม ดูโค้ดการแปลง Keras LSTM เป็นตัวอย่าง

แปลงเป็น LiteRT

ใช้ API TFLiteConverter.from_saved_model เพื่อแปลงเป็น LiteRT

กลไกภายใน

ตอนนี้เราจะอธิบายรายละเอียดระดับสูงของการออกแบบโดยรวมในการแปลงเป็น การดำเนินการที่ผสานรวมใน LiteRT

การเขียนการดำเนินการใน TensorFlow

การใช้ tf.function กับแอตทริบิวต์ฟังก์ชัน experimental_implements ช่วยให้ผู้ใช้สร้างการดำเนินการใหม่ได้อย่างชัดเจนโดยใช้ การดำเนินการดั้งเดิมของ TensorFlow และระบุอินเทอร์เฟซที่การดำเนินการแบบคอมโพสิต ที่เป็นผลลัพธ์จะใช้ ซึ่งมีประโยชน์อย่างมากเนื่องจากมีข้อมูลต่อไปนี้

  1. ขอบเขตที่กำหนดไว้อย่างดีสำหรับการดำเนินการแบบคอมโพสิตในกราฟ TensorFlow พื้นฐาน
  2. ระบุอินเทอร์เฟซที่การดำเนินการนี้ใช้โดยชัดแจ้ง อาร์กิวเมนต์ของ 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 ในตัวแปลงเพื่อรองรับกรณีการใช้งานการผสานการทำงานแบบคอมโพสิต โดยฟีเจอร์ใหม่ที่เพิ่มเข้ามามีดังนี้

  1. การนำเข้า โมเดลที่บันทึกไว้ของ TensorFlow ไปยัง MLIR
  2. fuse composite operations
  3. การวิเคราะห์ความสามารถในการเปลี่ยนแปลงของตัวแปร
  4. ตรึงตัวแปรแบบอ่านอย่างเดียวทั้งหมด

ซึ่งช่วยให้เราสามารถผสานการทำงานโดยใช้ฟังก์ชันที่แสดงการทำงานแบบคอมโพสิตก่อนที่จะมีการแทรกฟังก์ชันและตรึงตัวแปร

การใช้การผสานการทำงาน

มาดูรายละเอียดเพิ่มเติมเกี่ยวกับบัตรผ่านการผสานรวมการดำเนินงานกัน บัตรนี้จะทำสิ่งต่อไปนี้

  1. วนซ้ำฟังก์ชันทั้งหมดในโมดูล MLIR
  2. หากฟังก์ชันมีแอตทริบิวต์ tf._implements ระบบจะเรียกใช้ยูทิลิตีการผสานการทำงานที่เหมาะสมตามค่าแอตทริบิวต์
  3. ยูทิลิตีการผสานการทำงานจะทำงานกับตัวถูกดำเนินการและ แอตทริบิวต์ของฟังก์ชัน (ซึ่งทำหน้าที่เป็นอินเทอร์เฟซสำหรับการแปลง) และแทนที่ เนื้อหาของฟังก์ชันด้วยเนื้อหาฟังก์ชันที่เทียบเท่าซึ่งมีการ ดำเนินการที่ผสานรวม
  4. ในหลายกรณี เนื้อหาที่แทนที่จะมีการดำเนินการอื่นๆ นอกเหนือจากการ ดำเนินการที่ผสานรวม ซึ่งสอดคล้องกับการแปลงแบบคงที่บางอย่างในตัวถูกดำเนินการของฟังก์ชัน เพื่อให้ได้ตัวถูกดำเนินการของการดำเนินการที่ผสาน เนื่องจากการคำนวณเหล่านี้ทั้งหมดสามารถพับค่าคงที่ออกไปได้ จึงจะไม่มีอยู่ใน Flatbuffer ที่ส่งออกซึ่งจะมีเฉพาะการดำเนินการที่ผสานรวมเท่านั้น

ข้อมูลโค้ดจากบัตรที่แสดงเวิร์กโฟลว์หลักมีดังนี้

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

ต่อไปนี้คือข้อมูลโค้ดที่แสดงการแมปการดำเนินการแบบคอมโพสิตนี้กับการดำเนินการที่ผสานรวมใน LiteRT โดยใช้ประโยชน์จากฟังก์ชันเป็นอินเทอร์เฟซการแปลง

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