Chuyển đổi TensorFlow RNN sang TensorFlow Lite

Tổng quan

TensorFlow Lite hỗ trợ chuyển đổi các mô hình TensorFlow RNN sang các hoạt động LSTM hợp nhất của TensorFlow Lite. Các hoạt động hợp nhất tồn tại để tối đa hóa hiệu suất của việc triển khai hạt nhân cơ bản của chúng, cũng như cung cấp giao diện cấp cao hơn để xác định các phép biến đổi phức tạp như lượng tử hóa.

Vì có nhiều biến thể của API RNN trong TensorFlow nên cách tiếp cận của chúng tôi gồm hai phần:

  1. Cung cấp hỗ trợ riêng cho các API TensorFlow RNN tiêu chuẩn như Keras LSTM. Đây là tùy chọn được đề nghị.
  2. Cung cấp giao diện vào cơ sở hạ tầng chuyển đổi để triển khai RNN do người dùng xác định để cắm vào và chuyển đổi sang TensorFlow Lite. Chúng tôi cung cấp một số ví dụ sẵn có về chuyển đổi như vậy bằng cách sử dụng giao diện LSTMCellSimpleLayerNormalizedLSTMCellSimple RNN của lingvo.

API chuyển đổi

Tính năng này là một phần của bản phát hành TensorFlow 2.3. Nó cũng có sẵn thông qua tf-nightly pip hoặc từ head.

Chức năng chuyển đổi này khả dụng khi chuyển đổi sang TensorFlow Lite thông qua SavingModel hoặc trực tiếp từ mô hình Keras. Xem cách sử dụng ví dụ.

Từ mô hình đã lưu

# 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()

Từ mô hình Keras

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

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

Ví dụ

Keras LSTM to TensorFlow Lite Colab minh họa cách sử dụng từ đầu đến cuối với trình thông dịch TensorFlow Lite.

Hỗ trợ API TensorFlow RNNs

Chúng tôi hỗ trợ chuyển đổi ngay lập tức Keras LSTM sang TensorFlow Lite. Để biết chi tiết về cách thức hoạt động của tính năng này, vui lòng tham khảo giao diện Keras LSTM và logic chuyển đổi ở đây .

Điều quan trọng nữa là làm nổi bật hợp đồng LSTM của TensorFlow Lite liên quan đến định nghĩa hoạt động của Keras:

  1. Kích thước 0 của tensor đầu vào là kích thước lô.
  2. Kích thước 0 của tenxơ hồi quy_weight là số lượng đầu ra.
  3. Các tenxơ trọng số và nhân tái phát được hoán vị.
  4. Trọng số chuyển đổi, nhân hồi quy chuyển vị và các tensor phân cực được chia thành 4 tensor có kích thước bằng nhau dọc theo chiều 0. Chúng tương ứng với cổng đầu vào, cổng quên, ô và cổng đầu ra .

Các biến thể LSTM của Keras

chuyên ngành thời gian

Người dùng có thể chọn chuyên ngành thời gian hoặc không chuyên ngành thời gian. Keras LSTM thêm thuộc tính thời gian chính vào thuộc tính def của hàm. Đối với LSTM chuỗi một chiều, chúng ta có thể chỉ cần ánh xạ tới thuộc tính chính thời gian của unidirecional_sequence_lstm.

LSTM hai chiều

LSTM hai chiều có thể được triển khai với hai lớp LSTM của Keras, một lớp chuyển tiếp và một lớp lùi, xem ví dụ tại đây . Khi chúng tôi thấy thuộc tính go_backward, chúng tôi nhận ra đó là LSTM lùi, sau đó chúng tôi nhóm LSTM tiến và lùi lại với nhau. Đây là công việc trong tương lai. Hiện tại, điều này tạo ra hai hoạt động UnidirectionalSequenceLSTM trong mô hình TensorFlow Lite.

Ví dụ về chuyển đổi LSTM do người dùng xác định

TensorFlow Lite cũng cung cấp một cách để chuyển đổi việc triển khai LSTM do người dùng xác định. Ở đây chúng tôi sử dụng LSTM của Lingvo làm ví dụ về cách triển khai điều đó. Để biết chi tiết, vui lòng tham khảo giao diện lingvo.LSTMCellSimple và logic chuyển đổi tại đây . Chúng tôi cũng cung cấp một ví dụ cho một định nghĩa LSTM khác của Lingvo trong giao diện lingvo.LayerNormalizedLSTMCellSimple và logic chuyển đổi của nó tại đây .

“Mang TensorFlow RNN của riêng bạn” đến TensorFlow Lite

Nếu giao diện RNN của người dùng khác với giao diện được hỗ trợ tiêu chuẩn thì có một số tùy chọn:

Tùy chọn 1: Viết mã bộ điều hợp trong python TensorFlow để điều chỉnh giao diện RNN cho giao diện Keras RNN. Điều này có nghĩa là một tf.function với chú thích tf_implements trên hàm của giao diện RNN được tạo giống hệt với hàm được tạo bởi lớp Keras LSTM. Sau này, API chuyển đổi tương tự được sử dụng cho Keras LSTM sẽ hoạt động.

Tùy chọn 2: Nếu không thể thực hiện được điều trên (ví dụ: Keras LSTM thiếu một số chức năng hiện được cung cấp bởi chuẩn hóa lớp giống như LSTM hợp nhất của TensorFlow Lite), thì hãy mở rộng trình chuyển đổi TensorFlow Lite bằng cách viết mã chuyển đổi tùy chỉnh và cắm nó vào chuẩn bị -composite-functions MLIR-pass tại đây . Giao diện của hàm phải được xử lý giống như một hợp đồng API và phải chứa các đối số cần thiết để chuyển đổi sang các hoạt động LSTM TensorFlow Lite hợp nhất - tức là đầu vào, độ lệch, trọng số, phép chiếu, chuẩn hóa lớp, v.v. Tốt hơn là các tensor được truyền làm đối số cho điều này hàm có thứ hạng đã biết (tức là RankedTensorType trong MLIR). Điều này giúp việc viết mã chuyển đổi có thể coi các tensor này là RankedTensorType dễ dàng hơn nhiều và giúp chuyển đổi chúng thành tensor được xếp hạng tương ứng với các toán hạng của toán tử TensorFlow Lite hợp nhất.

Một ví dụ hoàn chỉnh về luồng chuyển đổi như vậy là chuyển đổi LSTMCellSimple sang TensorFlow Lite của Lingvo.

LSTMCellSimple trong Lingvo được định nghĩa ở đây . Các mô hình được đào tạo với ô LSTM này có thể được chuyển đổi sang TensorFlow Lite như sau:

  1. Bao bọc tất cả các cách sử dụng LSTMCellSimple trong một tf.function bằng chú thích tf_implements được gắn nhãn như vậy (ví dụ: lingvo.LSTMCellSimple sẽ là một tên chú thích phù hợp ở đây). Đảm bảo tf.function được tạo khớp với giao diện của hàm dự kiến ​​trong mã chuyển đổi. Đây là hợp đồng giữa tác giả mô hình thêm chú thích và mã chuyển đổi.
  2. Mở rộng thẻ chức năng chuẩn bị tổng hợp để cắm một op tổng hợp tùy chỉnh vào chuyển đổi op LSTM hợp nhất của TensorFlow Lite. Xem mã chuyển đổi LSTMCellSimple .

    Hợp đồng chuyển đổi:

  3. Các tensor trọng lượnghình chiếu được hoán vị.

  4. {Đầu vào, lặp lại} đến {ô, cổng đầu vào, cổng quên, cổng đầu ra} được trích xuất bằng cách cắt tenxơ trọng số chuyển đổi.

  5. {Độ lệch} đến {ô, cổng đầu vào, cổng quên, cổng đầu ra} được trích xuất bằng cách cắt tenxơ sai lệch.

  6. Phép chiếu được trích xuất bằng cách cắt tenxơ chiếu chuyển vị.

  7. Chuyển đổi tương tự được viết cho LayerNormalizedLSTMCellSimple .

  8. Phần còn lại của cơ sở hạ tầng chuyển đổi TensorFlow Lite, bao gồm tất cả các đường chuyền MLIR được xác định cũng như lần xuất cuối cùng sang bộ đệm phẳng TensorFlow Lite có thể được sử dụng lại.

Các vấn đề/hạn chế đã biết

  1. Hiện tại chỉ có hỗ trợ chuyển đổi Keras LSTM không trạng thái (hành vi mặc định trong Keras). Chuyển đổi Stateful Keras LSTM là công việc trong tương lai.
  2. Vẫn có thể mô hình hóa lớp Keras LSTM có trạng thái bằng cách sử dụng lớp Keras LSTM không trạng thái bên dưới và quản lý trạng thái một cách rõ ràng trong chương trình người dùng. Chương trình TensorFlow như vậy vẫn có thể được chuyển đổi thành TensorFlow Lite bằng tính năng được mô tả ở đây.
  3. LSTM hai chiều hiện được mô hình hóa thành hai hoạt động UnidirectionalSequenceLSTM trong TensorFlow Lite. Điều này sẽ được thay thế bằng một op BidirectionalSequenceLSTM duy nhất.