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

ภาพรวม

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

เนื่องจาก TensorFlow มี RNN API หลายรูปแบบ แนวทางของเราจึงมี 2 ส่วน ดังนี้

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

Converter API

ฟีเจอร์นี้เป็นส่วนหนึ่งของ TensorFlow เวอร์ชัน 2.3 นอกจากนี้ยังพร้อมใช้งานผ่าน pip ของ tf-nightly หรือจากส่วนหัว

ฟังก์ชันการทำงานของ Conversion นี้จะพร้อมใช้งานเมื่อแปลงเป็น LiteRT ผ่าน 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 ไปยัง LiteRT Colab แสดงให้เห็นการใช้งานตั้งแต่ต้นจนจบด้วยตัวแปล LiteRT

รองรับ TensorFlow RNNs APIs

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

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

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

LSTM เวอร์ชันต่างๆ ของ Keras

Time Major

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

LSTM แบบ 2 ทิศทาง

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

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

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

"นำ RNN ของ TensorFlow มาใช้เอง" ใน LiteRT

หากอินเทอร์เฟซ RNN ของผู้ใช้แตกต่างจากอินเทอร์เฟซมาตรฐานที่รองรับ คุณจะมีตัวเลือก 2 อย่างดังนี้

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

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

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

LSTMCellSimple ใน Lingvo มีคำจำกัดความที่นี่ คุณสามารถแปลงโมเดลที่ฝึกด้วยเซลล์ LSTM นี้เป็น LiteRT ได้ดังนี้

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

  2. ขยายการส่งผ่านฟังก์ชันเตรียมคอมโพสิตเพื่อเสียบการดำเนินการคอมโพสิตที่กำหนดเอง เพื่อการแปลงการดำเนินการ LSTM ที่ผสานรวม LiteRT ดูโค้ดการแปลง LSTMCellSimple

    สัญญาการแปลง

  3. มีการสลับเมตริกซ์ของเทนเซอร์น้ำหนักและการฉายภาพ

  4. ระบบจะดึงข้อมูล{input, recurrent} ไปยัง {cell, input gate, forget gate, output gate} โดยการแบ่งเทนเซอร์น้ำหนักที่เปลี่ยนแล้ว

  5. {bias} ของ {cell, input gate, forget gate, output gate} จะ ดึงข้อมูลโดยการแบ่งเทนเซอร์อคติ

  6. ระบบจะดึงการคาดการณ์ออกมาโดยการแบ่งเทนเซอร์การคาดการณ์ที่เปลี่ยนตำแหน่ง

  7. มีการเขียน Conversion ที่คล้ายกันสำหรับ LayerNormalizedLSTMCellSimple

  8. โครงสร้างพื้นฐาน Conversion ของ LiteRT ที่เหลือ รวมถึงการส่งผ่าน MLIR ทั้งหมดที่กำหนดไว้ ตลอดจนการส่งออกสุดท้ายไปยัง Flatbuffer ของ LiteRT สามารถนำกลับมาใช้ใหม่ได้

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

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