Chuyển đổi TensorFlow RNN thành TensorFlow Lite

Tổng quan

TensorFlow Lite hỗ trợ chuyển đổi các mô hình TensorFlow RNN sang các thao tác LSTM kết hợp của TensorFlow Lite. Các toán tử kết hợp tồn tại để tối đa hoá hiệu suất của việc triển khai hạt nhân cơ bản, cũng như cung cấp một 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ử hoá.

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 là gấp đôi:

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

API trình 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. Bạn cũng có thể sử dụng báo cáo này thông qua pip tf-nightly hoặc từ đầu.

Bạn có thể sử dụng chức năng chuyển đổi này khi chuyển đổi sang TensorFlow Lite thông qua mô hình đã lưu hoặc từ mô hình Keras trực tiếp. Hãy xem ví dụ về cách sử dụng.

Từ mẫu đã 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 đến TensorFlow Lite Colab minh hoạ quá trình sử dụng toàn diện bằng trình phiên dịch TensorFlow Lite.

Hỗ trợ các API TensorFlow RNNs

Chúng tôi hỗ trợ việc chuyển đổi trực tiếp Keras LSTM sang TensorFlow Lite. Để biết thông tin chi tiết về cách 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 tại đây.

Một điểm 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 tensor recurrent_weight là số đầu ra.
  3. Các tensor weightrecurrent_kernel được hoán vị.
  4. Trọng số hoán vị, tensor chuyển đổi recurrent_kernel và tensor bias được chia thành 4 tensor có kích thước bằng nhau dọc theo chiều 0. Các số này tương ứng với cổng đầu vào, cổng quên, ô và cổng đầu ra.

Biến thể LSTM của Keras

Thời gian chính

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

LSTM hai chiều

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

Ví dụ về lượt 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 các phương thức triển khai LSTM do người dùng xác định. Ở đây, chúng tôi dùng LSTM của Lingvo làm ví dụ về cách triển khai. Để biết thông tin 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 ví dụ cho một định nghĩa khác của Lingvo về LSTM trong giao diện lingvo.LayernormalizedLSTMCellSimple và logic chuyển đổi của giao diện đó tại đây.

"Sử dụng TensorFlow RNN của riêng bạn" cho TensorFlow Lite

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

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

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

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

Định nghĩa về LSTMCellSimple trong Lingvo tại đây. Bạn có thể chuyển đổi các mô hình được huấn luyện bằng ô LSTM này thành TensorFlow Lite như sau:

  1. Gói tất cả các trường hợp sử dụng LSTMCellSimple trong một tf.function bằng một chú giải tf_implementations được gắn nhãn như vậy (ví dụ: lingvo.LSTMCellSimple sẽ là một tên chú thích phù hợp tại đây). Hãy đảm bảo tf.hàm được tạo khớp với giao diện của hàm được mong đợi 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 việc truyền hàm tổng hợp chuẩn bị để cắm vào một thành phần kết hợp tuỳ chỉnh cho hoạt động chuyển đổi cơ chế hoạt động LSTM kết hợp của TensorFlow Lite. Hãy xem mã chuyển đổi LSTMCellSimple.

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

  3. Tensor Weight (Cân nặng) và projection (phép chiếu) được hoán vị.

  4. {input, recurrent} sang {cell, cổng vào, cổng quên, cổng đầu ra} được trích xuất bằng cách cắt lát tensor trọng số hoán vị.

  5. {bias} đến {cell, cổng nhập, cổng quên, cổng đầu ra} được trích xuất bằng cách cắt tensor xu hướng.

  6. Phép chiếu được trích xuất bằng cách cắt lát tensor chiếu hoán vị.

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

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

Vấn đề/giới hạn đã biết

  1. Hiện tại, chúng tôi chỉ hỗ trợ việc chuyển đổi Keras LSTM không có 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. Bạn vẫn có thể lập mô hình lớp Keras LSTM có trạng thái bằng cách sử dụng lớp Keras LSTM không có trạng thái cơ bản và quản lý trạng thái đó một cách rõ ràng trong chương trình người dùng. Bạn vẫn có thể chuyển đổi chương trình TensorFlow như vậy 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 hoá thành 2 thao tác UnidirectionalSequenceLSTM trong TensorFlow Lite. Tính năng này sẽ được thay thế bằng một hoạt động BidirectionalSequenceLSTM duy nhất.