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

ภาพรวม

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

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

การดำเนินการแบบหลอมรวมคืออะไร

การวาดภาพ

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

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

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

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

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

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

ความท้าทายกับการดำเนินงานแบบหลอมรวม

การแปลงการดำเนินการแบบผสมจาก TensorFlow ไปเป็นการดำเนินการแบบหลอมรวมใน TensorFlow Lite ถือเป็นปัญหาหนัก นี้เป็นเพราะ:

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

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

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

รวมการดำเนินการคอมโพสิตไว้ใน 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

ใช้การดำเนินการแบบหลอมรวมของคุณเป็นการดำเนินการ TFLite Custom - ดู คำแนะนำ

โปรดทราบว่าชื่อที่จะลงทะเบียน op ด้วยควรคล้ายกับชื่อที่ระบุในแอตทริบิวต์ name ในลายเซ็นนำไปใช้

ตัวอย่างสำหรับ op ในตัวอย่างคือ

  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 เป็นการดำเนินการแบบหลอมรวมของ TensorFlow Lite มีดังต่อไปนี้:

การวาดภาพ

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

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

เขียนโค้ด Conversion

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

ในพาสเตรียมคอมโพสิตฟังก์ชัน ให้ปลั๊กอิน โค้ด 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. หยุดตัวแปรอ่านอย่างเดียวทั้งหมด

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

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

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

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

นี่คือข้อมูลโค้ดจากบัตรผ่านที่แสดงขั้นตอนการทำงานหลัก:

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 โดยใช้ประโยชน์จากฟังก์ชันเป็นอินเทอร์เฟซการแปลง

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