Konwersja TensorFlow RNN do TensorFlow Lite

Przegląd

TensorFlow Lite obsługuje konwersję modeli TensorFlow RNN na połączone operacje LSTM TensorFlow Lite. Operacje fuzyjne mają na celu maksymalizację wydajności podstawowych implementacji jądra, a także zapewniają interfejs wyższego poziomu do definiowania złożonych transformacji, takich jak kwantyzacja.

Ponieważ w TensorFlow istnieje wiele wariantów interfejsów API RNN, nasze podejście było dwojakie:

  1. Zapewnij natywną obsługę standardowych interfejsów API TensorFlow RNN, takich jak Keras LSTM. Jest to zalecana opcja.
  2. Zapewnij interfejs do infrastruktury konwersji dla zdefiniowanych przez użytkownika implementacji RNN , aby podłączyć się i przekonwertować na TensorFlow Lite. Oferujemy kilka gotowych przykładów takiej konwersji przy użyciu interfejsów LSTMCellSimple i LayerNormalizedLSTMCellSimple RNN firmy lingvo.

API konwertera

Ta funkcja jest częścią wersji TensorFlow 2.3. Jest również dostępny poprzez tf-nightly pip lub z głowy.

Ta funkcja konwersji jest dostępna w przypadku konwersji do TensorFlow Lite za pośrednictwem SavedModel lub bezpośrednio z modelu Keras. Zobacz przykładowe zastosowania.

Z zapisanego modelu

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

Z modelu Keras

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

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

Przykład

Keras LSTM do TensorFlow Lite Colab ilustruje kompleksowe wykorzystanie interpretera TensorFlow Lite.

Obsługiwane interfejsy API TensorFlow RNN

Obsługujemy gotową do użycia konwersję Keras LSTM do TensorFlow Lite. Szczegółowe informacje o tym, jak to działa, można znaleźć w interfejsie Keras LSTM i do logiki konwersji tutaj .

Ważne jest również podkreślenie umowy LSTM TensorFlow Lite w odniesieniu do definicji operacji Keras:

  1. Wymiar 0 tensora wejściowego to wielkość partii.
  2. Wymiar 0 tensora recurrent_weight to liczba wyników.
  3. Tensory wagi i recurrent_kernel są transponowane.
  4. Transponowana waga, transponowane tensory recurrent_kernel i bias są podzielone na 4 tensory o równej wielkości wzdłuż wymiaru 0. Odpowiadają one bramce wejściowej, bramce zapominania, komórce i bramce wyjściowej .

Warianty Keras LSTM

Czas, major

Użytkownicy mogą wybrać czas główny lub brak czasu głównego. Keras LSTM dodaje atrybut czasowo-główny do atrybutów funkcji def. W przypadku sekwencji jednokierunkowej LSTM możemy po prostu odwzorować główny atrybut czasu unidirecional_sequence_lstm.

Dwukierunkowy LSTM

Dwukierunkowy LSTM można wdrożyć za pomocą dwóch warstw Keras LSTM, jednej do przodu i jednej do tyłu, zobacz przykłady tutaj . Gdy zobaczymy atrybut go_backward, rozpoznajemy go jako wsteczny LSTM, a następnie grupujemy razem LSTM do przodu i do tyłu. To jest praca na przyszłość. Obecnie tworzy to dwie operacje UnidirectSequenceLSTM w modelu TensorFlow Lite.

Zdefiniowane przez użytkownika przykłady konwersji LSTM

TensorFlow Lite zapewnia również sposób konwersji zdefiniowanych przez użytkownika implementacji LSTM. Tutaj używamy LSTM firmy Lingvo jako przykładu tego, jak można to wdrożyć. Aby uzyskać szczegółowe informacje, zapoznaj się z interfejsem lingvo.LSTMCellSimple i logiką konwersji tutaj . Podajemy również przykład innej definicji LSTM firmy Lingvo w interfejsie lingvo.LayerNormalizedLSTMCellSimple i logikę konwersji tutaj .

„Przynieś swój własny TensorFlow RNN” do TensorFlow Lite

Jeśli interfejs RNN użytkownika różni się od standardowo obsługiwanych, istnieje kilka opcji:

Opcja 1: Napisz kod adaptera w Pythonie TensorFlow, aby dostosować interfejs RNN do interfejsu Keras RNN. Oznacza to funkcję tf. z adnotacją tf_implements na wygenerowanej funkcji interfejsu RNN, która jest identyczna z funkcją wygenerowaną przez warstwę Keras LSTM. Następnie będzie działać ten sam interfejs API konwersji, który jest używany w Keras LSTM.

Opcja 2: Jeśli powyższe nie jest możliwe (np. w Keras LSTM brakuje niektórych funkcjonalności, które są obecnie udostępniane przez połączoną operację LSTM TensorFlow Lite w normalizacji warstw), następnie rozszerz konwerter TensorFlow Lite, pisząc niestandardowy kod konwersji i podłącz go do modułu przygotowania -composite-functions MLIR-przekaż tutaj . Interfejs funkcji powinien być traktowany jak kontrakt API i powinien zawierać argumenty potrzebne do konwersji na połączone operacje TensorFlow Lite LSTM - tj. dane wejściowe, obciążenie, wagi, projekcja, normalizacja warstw itp. Preferowane jest, aby tensory przekazywane były jako argumenty tej funkcji funkcja mająca znaną rangę (tj. RankingTensorType w MLIR). Ułatwia to znacznie pisanie kodu konwersji, który może przyjąć te tensory jako RankingTensorType i pomaga przekształcić je w tensory rankingowe odpowiadające połączonym operandom operatora TensorFlow Lite.

Kompletnym przykładem takiego przepływu konwersji jest konwersja LSTMCellSimple do TensorFlow Lite firmy Lingvo.

Tutaj zdefiniowano LSTMCellSimple w Lingvo. Modele przeszkolone za pomocą tej komórki LSTM można przekonwertować na TensorFlow Lite w następujący sposób:

  1. Zawiń wszystkie zastosowania LSTMCellSimple w funkcję tf. z adnotacją tf_implements oznaczoną jako taka (np. lingvo.LSTMCellSimple będzie tutaj dobrą nazwą adnotacji). Upewnij się, że wygenerowana funkcja tf.function odpowiada interfejsowi funkcji oczekiwanej w kodzie konwersji. Jest to umowa pomiędzy autorem modelu, który doda adnotację, a kodem konwersji.
  2. Rozszerz przepustkę funkcji przygotowania złożonego, aby podłączyć niestandardową opcję złożoną do konwersji operacji LSTM TensorFlow Lite. Zobacz kod konwersji LSTMCellSimple .

    Umowa przekształceniowa:

  3. Tensory wagi i projekcji są transponowane.

  4. Wartości {input, recurrent} do {komórka, bramka wejściowa, bramka zapominania, bramka wyjściowa} są wyodrębniane poprzez przecięcie transponowanego tensora wagi.

  5. {odchylenie} do {komórka, bramka wejściowa, bramka zapominania, bramka wyjściowa} są wyodrębniane poprzez przecięcie tensora odchylenia.

  6. Rzut jest wyodrębniany poprzez przecięcie transponowanego tensora projekcji.

  7. Podobna konwersja została napisana dla LayerNormalizedLSTMCellSimple .

  8. Pozostałą część infrastruktury konwersji TensorFlow Lite, w tym wszystkie zdefiniowane przebiegi MLIR , a także końcowy eksport do płaskiego bufora TensorFlow Lite, można ponownie wykorzystać.

Znane problemy/ograniczenia

  1. Obecnie obsługiwana jest tylko konwersja bezstanowego Keras LSTM (domyślne zachowanie w Keras). Konwersja stanowa Keras LSTM to praca na przyszłość.
  2. Nadal możliwe jest modelowanie stanowej warstwy Keras LSTM przy użyciu podstawowej bezstanowej warstwy Keras LSTM i jawne zarządzanie stanem w programie użytkownika. Taki program TensorFlow można nadal przekonwertować na TensorFlow Lite za pomocą opisanej tutaj funkcji.
  3. Dwukierunkowy LSTM jest obecnie modelowany jako dwie operacje UnilateralSequenceLSTM w TensorFlow Lite. Zostanie to zastąpione pojedynczą operacją BilateralSequenceLSTM.