Konwersja RNN z TensorFlow na LiteRT

Przegląd

LiteRT obsługuje konwertowanie modeli RNN TensorFlow na zintegrowane operacje LSTM LiteRT. Operacje scalone mają na celu maksymalizację wydajności implementacji jądra bazowego, a także zapewnienie interfejsu wyższego poziomu do definiowania złożonych przekształceń, takich jak kwantyzacja.

W TensorFlow jest wiele wariantów interfejsów API sieci RNN, dlatego nasze podejście jest dwutorowe:

  1. Zapewnia natywną obsługę standardowych interfejsów API TensorFlow RNN, takich jak Keras LSTM. Jest to opcja zalecana.
  2. Udostępnij interfejs infrastruktury konwersji, aby zdefiniowane przez użytkownika implementacje sieci RNN można było podłączyć i przekształcić w LiteRT. Podajemy kilka gotowych przykładów takiej konwersji z użyciem interfejsów RNN LSTMCellSimple i LayerNormalizedLSTMCellSimple w bibliotece Lingvo.

Converter API

Ta funkcja jest częścią wersji TensorFlow 2.3. Jest ona też dostępna w pakiecie tf-nightly lub w wersji głównej.

Ta funkcja konwersji jest dostępna podczas konwertowania na LiteRT za pomocą zapisanego modelu lub bezpośrednio z modelu Keras. Zobacz przykłady użycia.

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 to LiteRT Colab ilustruje kompleksowe użycie interpretera LiteRT.

Obsługiwane interfejsy API sieci RNN TensorFlow

Obsługujemy gotową konwersję modelu Keras LSTM na LiteRT. Szczegółowe informacje o tym, jak to działa, znajdziesz w interfejsie Keras LSTM oraz w logice konwersji tutaj.

Ważne jest również podkreślenie kontraktu LSTM LiteRT w odniesieniu do definicji operacji Keras:

  1. Wymiar 0 tensora wejściowego to rozmiar wsadu.
  2. Wymiar 0 tensora recurrent_weight to liczba danych wyjściowych.
  3. Tensory weight i recurrent_kernel są transponowane.
  4. Transponowane tensory wagi, transponowane tensory recurrent_kernel i bias są dzielone na 4 tensory o równych rozmiarach wzdłuż wymiaru 0. Odpowiadają one bramie wejściowej, bramie zapominania, komórce i bramie wyjściowej.

Warianty LSTM w Keras

Time major

Użytkownicy mogą wybrać opcję „time-major” lub „no time-major”. Keras LSTM dodaje atrybut time-major do atrybutów definicji funkcji. W przypadku jednokierunkowej sekwencji LSTM możemy po prostu zmapować atrybut time major jednokierunkowej sekwencji LSTM.

Dwukierunkowa sieć LSTM

Dwukierunkową sieć LSTM można zaimplementować za pomocą 2 warstw LSTM w Kerasie – jednej do przetwarzania sekwencji w kierunku do przodu i jednej do przetwarzania sekwencji w kierunku do tyłu. Przykłady znajdziesz tutaj. Gdy zobaczymy atrybut go_backward, rozpoznamy go jako sieć LSTM działającą w kierunku wstecznym, a następnie zgrupujemy sieci LSTM działające w kierunku do przodu i wstecznym. To jest praca na przyszłość. Obecnie powoduje to utworzenie w modelu LiteRT 2 operacji UnidirectionalSequenceLSTM.

Przykłady konwersji z użyciem zdefiniowanej przez użytkownika sieci LSTM

LiteRT umożliwia też konwertowanie zdefiniowanych przez użytkownika implementacji LSTM. Jako przykład implementacji tego rozwiązania użyjemy sieci LSTM z Lingvo. Szczegółowe informacje znajdziesz w interfejsie lingvo.LSTMCellSimple i logice konwersji tutaj. Przykład innej definicji LSTM w Lingvo znajdziesz w interfejsie lingvo.LayerNormalizedLSTMCellSimple i jego logice konwersji tutaj.

„Przenoszenie własnej sieci RNN TensorFlow” do LiteRT

Jeśli interfejs sieci RNN użytkownika różni się od standardowych obsługiwanych interfejsów, masz kilka możliwości:

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

Opcja 2: jeśli powyższe rozwiązanie nie jest możliwe (np.w Keras LSTM brakuje niektórych funkcji, które są obecnie udostępniane przez połączoną operację LSTM LiteRT, np. normalizacja warstwy), rozszerz konwerter LiteRT, pisząc niestandardowy kod konwersji i podłączając go do przekazywania MLIR prepare-composite-functions tutaj. Interfejs funkcji powinien być traktowany jak umowa interfejsu API i zawierać argumenty potrzebne do przekształcenia w zintegrowane operacje LiteRT LSTM, czyli dane wejściowe, odchylenie, wagi, projekcję, normalizację warstwy itp. Wskazane jest, aby tensory przekazywane jako argumenty do tej funkcji miały znany rangę (czyli RankedTensorType w MLIR). Ułatwia to pisanie kodu konwersji, który może przyjmować te tensory jako RankedTensorType i pomaga przekształcać je w tensory o określonej randze odpowiadające operandowi połączonego operatora LiteRT.

Kompletnym przykładem takiego przepływu konwersji jest konwersja z LSTMCellSimple w Lingvo na LiteRT.

Komórka LSTMCellSimple w Lingvo jest zdefiniowana tutaj. Modele wytrenowane za pomocą tej komórki LSTM można przekonwertować na LiteRT w ten sposób:

  1. Wszystkie użycia LSTMCellSimple należy umieścić w tf.function z adnotacją tf_implements oznaczoną w odpowiedni sposób (np. lingvo.LSTMCellSimple byłaby tutaj dobrą nazwą adnotacji). Upewnij się, że wygenerowana funkcja tf.function pasuje do interfejsu funkcji oczekiwanej w kodzie konwersji. Jest to umowa między autorem modelu dodającym adnotację a kodem konwersji.

  2. Rozszerz przejście prepare-composite-functions, aby włączyć niestandardową operację złożoną do konwersji operacji LSTM z fuzją LiteRT. Zobacz kod konwersji LSTMCellSimple.

    Umowa o konwersję:

  3. Tensory wagi i projekcji są transponowane.

  4. Wartości {input, recurrent} do {cell, input gate, forget gate, output gate} są wyodrębniane przez wycinanie transponowanego tensora wagi.

  5. Wartości {bias} do {cell, input gate, forget gate, output gate} są wyodrębniane przez wycinanie tensora odchylenia.

  6. Projekcja jest wyodrębniana przez wyodrębnienie transponowanego tensora projekcji.

  7. Podobna konwersja jest zapisywana dla funkcji LayerNormalizedLSTMCellSimple.

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

Znane problemy i ograniczenia

  1. Obecnie obsługiwane jest tylko przekształcanie bezstanowych sieci LSTM w Kerasie (domyślne działanie w Kerasie). Konwersja stanu sieci LSTM Keras to zadanie na przyszłość.
  2. Nadal można modelować warstwę LSTM z stanem w Keras, korzystając z bazowej warstwy LSTM bez stanu w Keras i zarządzając stanem bezpośrednio w programie użytkownika. Taki program TensorFlow można nadal przekonwertować na LiteRT za pomocą opisanej tutaj funkcji.
  3. Dwukierunkowy model LSTM jest obecnie modelowany jako 2 operacje UnidirectionalSequenceLSTM w LiteRT. Zostanie ona zastąpiona pojedynczym operatorem BidirectionalSequenceLSTM.