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

ภาพรวม

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

ตัวอย่างการใช้งานโครงสร้างพื้นฐานนี้คือ TensorFlow RNN ฟิวชันกับ LiteRT ตามรายละเอียดที่นี่

การดำเนินการรวมคืออะไร

การวาดภาพ

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

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

ประโยชน์ของการดำเนินการรวม

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

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

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

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

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

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

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

ด้วยเหตุนี้ การแปลงการดำเนินการแบบผสมผสานจึงได้รับการพิสูจน์แล้วว่าค่อนข้างท้าทาย

รวมการดำเนินการแบบผสมใน tf.function

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

ตัวอย่างเช่น

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 ระบุเป็นนัยนี้อยู่แล้ว

ใช้การดำเนินการที่กำหนดเองและลงทะเบียนกับล่าม TFLite

ใช้งานการดำเนินการ Fused ของคุณเป็นการดำเนินการ TFLite Custom - ดู วิธีการ

โปรดทราบว่าชื่อที่ใช้จดทะเบียนผู้ดำเนินการควรคล้ายกับชื่อ ที่ระบุในแอตทริบิวต์ 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);

กำลังแปลงจากแบบผสมเป็นการดำเนินการ Fused (ขั้นสูง)

สถาปัตยกรรมโดยรวมสำหรับการแปลงการดำเนินการแบบผสมของ TensorFlow เป็น การดำเนินการ Fused ของ LiteRT ได้ที่ด้านล่าง

การวาดภาพ

รวมการดำเนินการแบบผสมใน tf.function

ในซอร์สโค้ดของโมเดล TensorFlow ให้ระบุและแยกองค์ประกอบออกมา ใน tf.function พร้อมฟังก์ชัน experimental_implements ฟังก์ชันพิเศษ ดูตัวอย่างการค้นหาแบบฝัง จะกำหนดอินเทอร์เฟซและควรใช้อาร์กิวเมนต์เพื่อปรับใช้ Conversion Logic

เขียนโค้ด Conversion

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

ในคำสั่งเตรียมการฟังก์ชันคอมโพสิต ให้ปลั๊กอินใน Conversion รหัส

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

แปลงเป็น LiteRT

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

กลไกภายใน

ตอนนี้เราได้อธิบายรายละเอียดในระดับสูงเกี่ยวกับการออกแบบโดยรวมในการแปลงเป็น Fused ใน 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. การดำเนินการประกอบฟิวส์
  3. การวิเคราะห์การเปลี่ยนแปลงของตัวแปร
  4. ตรึงตัวแปรแบบอ่านอย่างเดียวทั้งหมด

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

การใช้การผสมผสานการดำเนินการ

มาดูรายละเอียดเพิ่มเติมของ Operation Fusion Pass กัน บัตรใบนี้จะทำสิ่งต่อไปนี้ ดังต่อไปนี้:

  1. วนซ้ำฟังก์ชันทั้งหมดในโมดูล MLIR
  2. หากฟังก์ชันมีแอตทริบิวต์ tf._APPLYs โดยอิงตามแอตทริบิวต์ จะเรียกยูทิลิตีฟิวชันการดำเนินการที่เหมาะสม
  3. ยูทิลิตีฟิวชันการดำเนินการจะทำงานกับตัวถูกดำเนินการของฟังก์ชันและ (ซึ่งทำหน้าที่เป็นอินเทอร์เฟซสำหรับ Conversion) และแทนที่ ส่วนเนื้อหาของฟังก์ชันซึ่งมีส่วนเนื้อหาของฟังก์ชันที่เทียบเท่าซึ่งมีค่า ที่ผสมผสานเข้าด้วยกัน
  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 */
}

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

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