การแปลง TensorFlow RNN เป็น TensorFlow Lite

ภาพรวม

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

เนื่องจากมี RNN API หลากหลายรูปแบบใน TensorFlow แนวทางของเราจึงเป็นสองเท่า:

  1. ให้ การสนับสนุนดั้งเดิมสำหรับ TensorFlow RNN API มาตรฐาน เช่น Keras LSTM นี่คือตัวเลือกที่แนะนำ
  2. จัดเตรียม อินเทอร์เฟซ ในโครงสร้างพื้นฐานการแปลงสำหรับ การใช้งาน RNN ที่ผู้ใช้กำหนด เพื่อเสียบปลั๊กและแปลงเป็น TensorFlow Lite เรามีตัวอย่างนอกกรอบสองสามตัวอย่างสำหรับการแปลงดังกล่าวโดยใช้อินเทอร์เฟซ LSTMCellSimple และ LayerNormalizedLSTMCellSimple RNN ของ lingvo

API ตัวแปลง

คุณลักษณะนี้เป็นส่วนหนึ่งของการเปิดตัว TensorFlow 2.3 นอกจากนี้ยังสามารถใช้ได้ผ่านทาง pip ของ tf-nightly หรือจากหัว

ฟังก์ชันการแปลงนี้ใช้งานได้เมื่อแปลงเป็น TensorFlow Lite ผ่าน SavedModel หรือจากโมเดล Keras โดยตรง ดูตัวอย่างการใช้งาน

จากโมเดลที่บันทึกไว้

# build a saved model. Here concrete_function is the exported function
# corresponding to the TensorFlow model containing one or more
# Keras LSTM layers.
saved_model, saved_model_dir = build_saved_model_lstm(...)
saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)

# Convert the model.
converter = TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

จากโมเดล Keras

# build a Keras model
keras_model = build_keras_lstm(...)

# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()

ตัวอย่าง

Keras LSTM ถึง TensorFlow Lite Colab แสดงให้เห็นการใช้งานตั้งแต่ต้นจนจบด้วยล่าม TensorFlow Lite

รองรับ API ของ TensorFlow RNN

เรารองรับการแปลง Keras LSTM เป็น TensorFlow Lite ได้ทันที สำหรับรายละเอียดเกี่ยวกับวิธีการทำงาน โปรดดู อินเทอร์เฟซ Keras LSTM และตรรกะการแปลง ที่นี่

สิ่งสำคัญอีกอย่างคือการเน้นสัญญา LSTM ของ TensorFlow Lite ที่เกี่ยวข้องกับคำจำกัดความการดำเนินการ Keras:

  1. มิติ 0 ของเทนเซอร์ อินพุต คือขนาดแบตช์
  2. มิติ 0 ของเทนเซอร์ recurrent_weight คือจำนวนเอาต์พุต
  3. เทนเซอร์ น้ำหนัก และเคอร์เนล recurrent_kernel ถูกย้าย
  4. น้ำหนักทรานสโพส, recurrent_kernel ที่ถูกย้าย และ ไบอัส เทนเซอร์จะถูกแบ่งออกเป็นเทนเซอร์ขนาดเท่ากัน 4 ตัวตามมิติ 0 ซึ่งสอดคล้องกับ เกตอินพุต, เกตลืม, เซลล์ และเกตเอาท์พุต

ตัวแปร Keras LSTM

เวลาที่สำคัญ

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

LSTM แบบสองทิศทาง

LSTM แบบสองทิศทางสามารถนำมาใช้กับเลเยอร์ Keras LSTM สองชั้น หนึ่งเลเยอร์สำหรับการส่งต่อและอีกเลเยอร์หนึ่งสำหรับการย้อนกลับ ดูตัวอย่าง ที่นี่ เมื่อเราเห็นแอตทริบิวต์ go_backward เราจะจำได้ว่ามันเป็น LSTM แบบย้อนกลับ จากนั้นเราจะจัดกลุ่ม LSTM แบบไปข้างหน้าและข้างหลังเข้าด้วยกัน นี่คืองานในอนาคต ในปัจจุบัน สิ่งนี้จะสร้างการดำเนินการ UniDirectionalSequenceLSTM สองรายการในโมเดล TensorFlow Lite

ตัวอย่างการแปลง LSTM ที่ผู้ใช้กำหนด

TensorFlow Lite ยังมีวิธีแปลงการใช้งาน LSTM ที่ผู้ใช้กำหนดอีกด้วย ที่นี่เราใช้ LSTM ของ Lingvo เป็นตัวอย่างของวิธีการนำไปปฏิบัติ สำหรับรายละเอียด โปรดดู อินเทอร์เฟซ lingvo.LSTMCellSimple และตรรกะการแปลง ที่นี่ นอกจากนี้เรายังมีตัวอย่างคำจำกัดความ LSTM อื่นๆ ของ Lingvo ใน อินเทอร์เฟซ lingvo.LayerNormalizedLSTMCellSimple และตรรกะการแปลงได้ ที่นี่

“นำ TensorFlow RNN ของคุณเอง” มาสู่ TensorFlow Lite

หากอินเทอร์เฟซ RNN ของผู้ใช้แตกต่างจากอินเทอร์เฟซมาตรฐานที่รองรับ มีสองตัวเลือก:

ตัวเลือกที่ 1: เขียนโค้ดอะแดปเตอร์ใน TensorFlow python เพื่อปรับอินเทอร์เฟซ RNN ให้เป็นอินเทอร์เฟซ Keras RNN นี่หมายถึง tf.function พร้อมด้วย คำอธิบายประกอบ tf_implements บนฟังก์ชันของอินเทอร์เฟซ RNN ที่สร้างขึ้น ซึ่งเหมือนกับฟังก์ชันที่สร้างโดยเลเยอร์ Keras LSTM หลังจากนี้ API การแปลงเดียวกันกับที่ใช้สำหรับ Keras LSTM จะใช้งานได้

ตัวเลือกที่ 2: หากสิ่งที่กล่าวมาข้างต้นเป็นไปไม่ได้ (เช่น Keras LSTM ขาดฟังก์ชันการทำงานบางอย่างที่ถูกเปิดเผยโดย LSTM op ที่หลอมรวมของ LSTM op ของ TensorFlow Lite เช่น การทำให้เลเยอร์ปกติ) ให้ขยายตัวแปลง TensorFlow Lite โดยการเขียนโค้ดการแปลงแบบกำหนดเองและเสียบเข้ากับการเตรียม -composite-functions MLIR-pass ที่นี่ อินเทอร์เฟซของฟังก์ชันควรได้รับการปฏิบัติเหมือนสัญญา API และควรมีอาร์กิวเมนต์ที่จำเป็นในการแปลงเป็นการดำเนินการ TensorFlow Lite LSTM แบบหลอมรวม เช่น อินพุต อคติ น้ำหนัก การฉายภาพ การทำให้เลเยอร์เป็นมาตรฐาน ฯลฯ จะดีกว่าสำหรับเทนเซอร์ที่ส่งผ่านเป็นอาร์กิวเมนต์ของสิ่งนี้ ฟังก์ชั่นเพื่อให้ทราบอันดับ (เช่น RankedTensorType ใน MLIR) ซึ่งช่วยให้เขียนโค้ดการแปลงที่สามารถถือว่าเทนเซอร์เหล่านี้เป็น RankedTensorType ได้ง่ายขึ้นมาก และช่วยแปลงให้เป็นเทนเซอร์ที่ได้รับการจัดอันดับที่สอดคล้องกับตัวถูกดำเนินการของตัวดำเนินการ TensorFlow Lite ที่หลอมรวม

ตัวอย่างที่สมบูรณ์ของโฟลว์การแปลงดังกล่าวคือการแปลง LSTMCellSimple เป็น TensorFlow Lite ของ Lingvo

LSTMCellSimple ใน Lingvo มีการกำหนดไว้ ที่นี่ โมเดลที่ได้รับการฝึกด้วยเซลล์ LSTM นี้สามารถแปลงเป็น TensorFlow Lite ได้ดังนี้:

  1. รวมการใช้งาน LSTMCellSimple ทั้งหมดไว้ใน tf.function ด้วยคำอธิบายประกอบ tf_implements ที่มีป้ายกำกับเช่นนี้ (เช่น lingvo.LSTMCellSimple จะเป็นชื่อคำอธิบายประกอบที่ดีที่นี่) ตรวจสอบให้แน่ใจว่า tf.function ที่สร้างขึ้นตรงกับอินเทอร์เฟซของฟังก์ชันที่คาดไว้ในโค้ด Conversion นี่เป็นสัญญาระหว่างผู้สร้างโมเดลในการเพิ่มคำอธิบายประกอบและโค้ดการแปลง
  2. ขยายการจัดเตรียมฟังก์ชันคอมโพสิตเพื่อเสียบคอมโพสิต op แบบกำหนดเองเข้ากับการแปลง op ของ TensorFlow Lite ที่หลอมรวม LSTM ดูโค้ดการแปลง LSTMCellSimple

    สัญญาแปลงสภาพ:

  3. เทนเซอร์ น้ำหนัก และ การฉายภาพ ถูกย้าย

  4. {input, recurrent} ไปยัง {cell, input gate,forget gate,output gate} จะถูกแยกออกมาโดยการแบ่งส่วนตุ้มน้ำหนักที่ย้ายออก

  5. {อคติ} ถึง {เซลล์, ประตูอินพุต, ลืมเกต, ประตูเอาท์พุต} ถูกแยกออกมาโดยการแบ่งไบแอสเทนเซอร์

  6. การฉายภาพ จะถูกดึงออกมาโดยการตัดเทนเซอร์การฉายภาพแบบทรานสโพส

  7. การแปลงที่คล้ายกันนี้เขียนขึ้นสำหรับ LayerNormalizedLSTMCellSimple

  8. โครงสร้างพื้นฐานการแปลง TensorFlow Lite ที่เหลือ รวมถึง MLIR Pass ทั้งหมดที่กำหนดไว้ รวมถึงการส่งออกขั้นสุดท้ายไปยัง Flatbuffer ของ TensorFlow Lite สามารถนำมาใช้ซ้ำได้

ปัญหา/ข้อจำกัดที่ทราบ

  1. ขณะนี้มีการรองรับเฉพาะการแปลง Keras LSTM ที่ไม่มีสถานะ (พฤติกรรมเริ่มต้นใน Keras) การแปลง Keras LSTM แบบมีสถานะเป็นงานในอนาคต
  2. ยังคงสามารถสร้างแบบจำลองเลเยอร์ Keras LSTM แบบมีสถานะได้โดยใช้เลเยอร์ Keras LSTM แบบไม่ระบุสถานะและจัดการสถานะอย่างชัดเจนในโปรแกรมผู้ใช้ โปรแกรม TensorFlow ดังกล่าวยังสามารถแปลงเป็น TensorFlow Lite ได้โดยใช้คุณสมบัติที่อธิบายไว้ที่นี่
  3. ปัจจุบัน LSTM แบบสองทิศทางได้รับการจำลองเป็นการดำเนินการ UniDirectionalSequenceLSTM สองรายการใน TensorFlow Lite สิ่งนี้จะถูกแทนที่ด้วย BiDirectionalSequenceLSTM op เดียว