ขอขอบคุณที่เข้าร่วม Google I/O ดูเซสชั่นทั้งหมดตามความต้องการ ดูตามความต้องการ

ฟิวชั่นการทำงานของ TensorFlow, ฟิวชั่นการทำงานของ TensorFlow

ภาพรวม

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

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

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

การวาดภาพ

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

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

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

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

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

TensorFlow Lite มีการดำเนินการแบบหลอมรวมหลายครั้งด้วยเหตุผลดังที่กล่าวไว้ข้างต้น การดำเนินการหลอมรวมเหล่านี้มักจะสอดคล้องกับการดำเนินการแบบผสมในโปรแกรม TensorFlow ต้นทาง ตัวอย่างของการดำเนินการแบบผสมใน TensorFlow ที่นำมาใช้เป็นการดำเนินการแบบหลอมรวมครั้งเดียวใน TensorFlow Lite รวมถึงการดำเนินการ RNN ต่างๆ เช่น LSTM ลำดับแบบทิศทางเดียวและแบบสองทิศทาง การบิดเบี้ยว (conv2d, การเพิ่มอคติ, relu) การเชื่อมต่ออย่างสมบูรณ์ (matmul, การเพิ่มอคติ, 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 coder.
    return kTfLiteOk;
  };
  reg.builtin_code = kTfLiteCustom;
  resolver->AddCustom(kOpName, &reg);

การแปลงจากการทำงานแบบผสมเป็นการทำงานแบบหลอมรวม (ขั้นสูง)

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

การวาดภาพ

ห่อการดำเนินการคอมโพสิตใน tf.function

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

เขียนโค้ดแปลง

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

ในการผ่าน prepare-composite-functions ให้เสียบปลั๊กอินใน โค้ดการแปลง ของคุณ

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

แปลงเป็น 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. ตรึงตัวแปรแบบอ่านอย่างเดียวทั้งหมด

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

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

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

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

นี่คือข้อมูลโค้ดที่แสดงการจับคู่การดำเนินการแบบผสมนี้กับการดำเนินการที่หลอมรวมใน 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());
  }