TensorFlow RNN을 LiteRT로 변환

개요

LiteRT는 TensorFlow RNN 모델을 LiteRT의 융합 LSTM 작업으로 변환하는 것을 지원합니다. 융합 작업은 기본 커널 구현의 성능을 최대화하고 양자화와 같은 복잡한 변환을 정의하는 상위 수준 인터페이스를 제공하기 위해 존재합니다.

TensorFlow에는 다양한 RNN API가 있으므로 접근 방식은 다음과 같이 두 가지입니다.

  1. Keras LSTM과 같은 표준 TensorFlow RNN API에 대한 네이티브 지원을 제공합니다. 권장되는 옵션입니다.
  2. 사용자 정의 RNN 구현을 연결하고 LiteRT로 변환할 수 있는 변환 인프라로의 인터페이스를 제공합니다. lingvo의 LSTMCellSimple 및 LayerNormalizedLSTMCellSimple RNN 인터페이스를 사용하여 이러한 변환의 기본 제공 예시를 제공합니다.

Converter API

이 기능은 TensorFlow 2.3 출시의 일부입니다. tf-nightly pip를 통해 또는 헤드에서 사용할 수도 있습니다.

이 변환 기능은 SavedModel을 통해 또는 Keras 모델에서 직접 LiteRT로 변환할 때 사용할 수 있습니다. 사용 예시를 참고하세요.

저장된 모델에서

# 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 to LiteRT Colab은 LiteRT 인터프리터의 엔드 투 엔드 사용을 보여줍니다.

지원되는 TensorFlow RNN API

Keras LSTM을 LiteRT로 즉시 변환하는 기능이 지원됩니다. 작동 방식에 관한 자세한 내용은 Keras LSTM 인터페이스 와 변환 로직(여기)을 참고하세요.

또한 Keras 작업 정의와 관련하여 LiteRT의 LSTM 계약을 강조하는 것도 중요합니다.

  1. 입력 텐서의 0번째 차원은 배치 크기입니다.
  2. recurrent_weight 텐서의 차원 0은 출력 수입니다.
  3. weight 및 recurrent_kernel 텐서가 전치됩니다.
  4. 전치된 가중치, 전치된 recurrent_kernel, bias 텐서는 차원 0을 따라 크기가 동일한 텐서 4개로 분할됩니다. 이는 입력 게이트, 망각 게이트, 셀, 출력 게이트에 해당합니다.

Keras LSTM 변형

시간 주요

사용자는 시간 중심 또는 시간 중심 아님을 선택할 수 있습니다. Keras LSTM은 함수 정의 속성에 시간 우선 속성을 추가합니다. 단방향 시퀀스 LSTM의 경우 unidirecional_sequence_lstm의 time major attribute에 간단히 매핑할 수 있습니다.

양방향 LSTM

양방향 LSTM은 두 개의 Keras LSTM 레이어로 구현할 수 있습니다. 하나는 순방향이고 하나는 역방향입니다. 여기에서 예를 참고하세요. go_backward 속성이 확인되면 이를 역방향 LSTM으로 인식하고 순방향 및 역방향 LSTM을 함께 그룹화합니다. 향후 작업입니다. 현재 이로 인해 LiteRT 모델에 UnidirectionalSequenceLSTM 작업이 두 개 생성됩니다.

사용자 정의 LSTM 변환 예

LiteRT는 사용자 정의 LSTM 구현을 변환하는 방법도 제공합니다. 여기서는 Lingvo의 LSTM을 이러한 구현의 예로 사용합니다. 자세한 내용은 lingvo.LSTMCellSimple 인터페이스와 여기의 변환 로직을 참고하세요. 또한 lingvo.LayerNormalizedLSTMCellSimple 인터페이스의 다른 Lingvo LSTM 정의와 변환 로직의 예도 여기에 제공됩니다.

LiteRT에 '자체 TensorFlow RNN 가져오기'

사용자의 RNN 인터페이스가 지원되는 표준 인터페이스와 다른 경우 다음 두 가지 옵션이 있습니다.

옵션 1: TensorFlow Python에서 어댑터 코드를 작성하여 RNN 인터페이스를 Keras RNN 인터페이스에 적응시킵니다. 이는 생성된 RNN 인터페이스의 함수에 tf_implements 주석이 있는 tf.function이 Keras LSTM 레이어에서 생성된 함수와 동일함을 의미합니다. 이후에는 Keras LSTM에 사용된 것과 동일한 변환 API가 작동합니다.

옵션 2: 위의 방법이 불가능한 경우 (예: Keras LSTM에 현재 LiteRT의 융합 LSTM 작업에서 노출되는 레이어 정규화와 같은 기능이 누락됨) 맞춤 변환 코드를 작성하여 LiteRT 변환기를 확장하고 여기의 prepare-composite-functions MLIR 패스에 연결합니다. 함수의 인터페이스는 API 계약처럼 취급해야 하며 융합된 LiteRT LSTM 작업으로 변환하는 데 필요한 인수(예: 입력, 편향, 가중치, 투영, 레이어 정규화 등)를 포함해야 합니다. 이 함수에 인수로 전달되는 텐서의 순위가 알려져 있는 것이 좋습니다(예: MLIR의 RankedTensorType). 이렇게 하면 이러한 텐서를 RankedTensorType으로 가정할 수 있는 변환 코드를 훨씬 쉽게 작성할 수 있으며, 융합된 LiteRT 연산자의 피연산자에 해당하는 순위가 지정된 텐서로 변환하는 데 도움이 됩니다.

이러한 변환 흐름의 전체 예는 Lingvo의 LSTMCellSimple에서 LiteRT로의 변환입니다.

Lingvo의 LSTMCellSimple은 여기에 정의되어 있습니다. 이 LSTM 셀로 학습된 모델은 다음과 같이 LiteRT로 변환할 수 있습니다.

  1. LSTMCellSimple의 모든 사용을 tf_implements 주석이 있는 tf.function으로 래핑합니다 (예: lingvo.LSTMCellSimple이 여기에서 적절한 주석 이름임). 생성된 tf.function이 변환 코드에서 예상되는 함수의 인터페이스와 일치하는지 확인합니다. 이는 주석과 전환 코드를 추가하는 모델 작성자 간의 계약입니다.

  2. 맞춤 복합 작업을 LiteRT 융합 LSTM 작업 변환에 연결하도록 prepare-composite-functions 패스를 확장합니다. LSTMCellSimple 변환 코드를 참고하세요.

    전환 계약:

  3. 가중치 및 투영 텐서가 전치됩니다.

  4. {cell, input gate, forget gate, output gate}에 대한 {input, recurrent}는 전치된 가중치 텐서를 슬라이싱하여 추출됩니다.

  5. {bias}에서 {cell, input gate, forget gate, output gate}은 편향 텐서를 슬라이싱하여 추출됩니다.

  6. 투영은 전치된 투영 텐서를 슬라이싱하여 추출됩니다.

  7. LayerNormalizedLSTMCellSimple에 유사한 변환이 작성됩니다.

  8. 정의된 모든 MLIR 패스와 LiteRT flatbuffer로의 최종 내보내기를 비롯한 나머지 LiteRT 변환 인프라는 재사용할 수 있습니다.

알려진 문제/제한사항

  1. 현재는 상태 비저장 Keras LSTM (Keras의 기본 동작) 변환만 지원됩니다. 상태 저장 Keras LSTM 변환은 향후 작업입니다.
  2. 기본 스테이트리스 Keras LSTM 레이어를 사용하고 사용자 프로그램에서 상태를 명시적으로 관리하여 스테이트풀 Keras LSTM 레이어를 모델링할 수는 있습니다. 이러한 TensorFlow 프로그램은 여기에 설명된 기능을 사용하여 LiteRT로 변환할 수 있습니다.
  3. 양방향 LSTM은 현재 LiteRT에서 두 개의 UnidirectionalSequenceLSTM 작업으로 모델링됩니다. 이는 단일 BidirectionalSequenceLSTM 작업으로 대체됩니다.