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

ภาพรวม

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

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

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

การวาดภาพ

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

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

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

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

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

TensorFlow Lite มีอินสแตนซ์ของการดำเนินการรวมกันหลายอินสแตนซ์ด้วยเหตุผลที่กล่าวไว้ข้างต้น โดยทั่วไปแล้ว การดําเนินการ Fused เหล่านี้จะสอดคล้องกับการดำเนินการผสมในโปรแกรม TensorFlow ต้นทาง ตัวอย่างการดำเนินการผสมใน TensorFlow ที่นำมาใช้เป็นการดำเนินการผสมเดี่ยวใน TensorFlow Lite รวมถึงการดำเนินการ RNN ต่างๆ เช่น ลำดับแบบ Unidirectional และ 2 Directional อย่าง LSTM, Convolution (conv2d, การให้น้ำหนักพิเศษ, relu), การเชื่อมต่อโดยสมบูรณ์ (Mattul, การให้น้ำหนักพิเศษ การเพิ่ม, Relu) และอื่นๆ ใน TensorFlow Lite ปัจจุบันการหาปริมาณของ LSTM จะใช้งานในการดำเนินการ LSTM ที่ Fused เท่านั้น

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

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

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

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

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

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

ในหลายกรณี บางส่วนของโมเดลจะแมปกับการดำเนินการเดียวใน TFLite ได้ วิธีนี้จะช่วยเรื่องประสิทธิภาพเมื่อเขียนการติดตั้งใช้งานที่เพิ่มประสิทธิภาพเพื่อการดำเนินการที่เฉพาะเจาะจง หากต้องการสร้างการดำเนินการ Fused ใน 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 บอกเป็นนัยถึงเรื่องนี้แล้ว

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

ติดตั้งใช้งานการดำเนินการ Fused เป็นการดำเนินการที่กำหนดเองของ TFLite โปรดดูinstructions

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

การแปลงจาก Composite เป็น Fused Fused (ขั้นสูง)

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

การวาดภาพ

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

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

เขียนโค้ด Conversion

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

ใน Pre-composite-Function ให้ปลั๊กอินในโค้ด Conversion

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

แปลงเป็น TensorFlow Lite

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

ขั้นสูง

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

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

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

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

ดูตัวอย่างจากตัวอย่างการดำเนินการแบบผสมที่กำหนดให้ใช้การค้นหาการฝัง การดำเนินการนี้จะแมปไปยังการดำเนินการที่ผสานรวมใน 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

การทําให้โมเดลใช้การดําเนินการแบบผสมผ่าน tf.function ตามที่แสดงให้เห็นด้านบน ทำให้สามารถสร้างโครงสร้างพื้นฐานทั่วไปเพื่อระบุและแปลงการดำเนินการดังกล่าวเป็นการหลอมรวมการดำเนินการ TensorFlow Lite ได้

การขยายตัวแปลง TensorFlow Lite

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

ในการใช้ประโยชน์จาก tf.function ที่มีฟีเจอร์ experimental_implements ระหว่างกระบวนการแปลง คุณจะต้องเก็บฟังก์ชันเหล่านี้ไว้จนกว่าจะถึงช่วงท้ายของกระบวนการแปลง

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

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

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

การใช้การดำเนินการเชื่อมโยง

มาดูรายละเอียดเพิ่มเติมของ Ops Fusion Pass กัน บัตรใบนี้จะมีคุณสมบัติดังนี้

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

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

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

นี่คือข้อมูลโค้ดที่แสดงการแมปการดำเนินการผสมนี้กับการดำเนินการที่รวมกันใน TensorFlow Lite ที่ใช้ประโยชน์จากฟังก์ชันเป็นอินเทอร์เฟซ 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());
  }