TensorFlow Lite로 TensorFlow RNN 변환

개요

TensorFlow Lite는 TensorFlow RNN 모델을 TensorFlow Lite의 통합된 LSTM 작업으로 변환할 수 있도록 지원합니다. 융합 연산은 기본 커널 구현의 성능을 극대화할 뿐만 아니라 양자화와 같은 복잡한 변환을 정의하기 위한 더 높은 수준의 인터페이스를 제공하기 위해 존재합니다.

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

  1. Keras LSTM과 같은 표준 TensorFlow RNN API를 기본적으로 지원합니다. 권장되는 옵션입니다.
  2. TensorFlow Lite로 플러그인하고 변환할 수 있도록 사용자 정의 RNN 구현을 위한 인터페이스를 변환 인프라에 제공합니다. Google에서는 링보의 LSTMCellSimple 및 LayerNormalizedLSTMCellSimple RNN 인터페이스를 사용하여 이러한 변환의 즉시 사용 가능한 몇 가지 예를 제공합니다.

변환기 API

이 기능은 TensorFlow 2.3 출시에 포함되어 있습니다. tf-nightly pip 또는 head를 통해 사용할 수도 있습니다.

이 변환 기능은 저장된 모델을 통해 또는 Keras 모델에서 직접 TensorFlow Lite로 변환할 때 사용할 수 있습니다. 사용 예 보기

저장된 모델에서

# 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 TensorFlow Lite Colab은 TensorFlow Lite 인터프리터를 사용한 엔드 투 엔드 사용을 보여줍니다.

지원되는 TensorFlow RNNs API

Keras LSTM을 TensorFlow Lite로 즉시 변환하는 기능을 지원합니다. 작동 방식에 관한 자세한 내용은 Keras LSTM 인터페이스 및 여기에서 변환 로직을 참고하세요.

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

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

Keras LSTM 변형

시간장조

사용자는 '주요 시간'을 선택하거나 '시간 시간 없음'을 선택할 수 있습니다. Keras LSTM은 함수 정의 속성에 time-major 속성을 추가합니다. 단방향 시퀀스 LSTM의 경우 단순히 unidirecional_sequence_lstm의 time 주요 속성에 매핑할 수 있습니다.

양방향 LSTM

양방향 LSTM은 두 개의 Keras LSTM 레이어로 구현할 수 있습니다. 하나는 정방향용이고 다른 하나는 역방향용입니다. 여기에서 예를 확인하세요. go_backward 속성이 표시되면 뒤로 LSTM으로 인식된 다음 정방향 및 역방향 LSTM을 함께 그룹화합니다. 이는 향후 작업입니다. 현재 이렇게 하면 TensorFlow Lite 모델에 단방향 SequenceLSTM 작업 두 개가 생성됩니다.

사용자 정의 LSTM 변환 예

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

TensorFlow Lite로 '자체 TensorFlow RNN 가져오기'

사용자의 RNN 인터페이스가 표준 지원 인터페이스와 다른 경우 다음과 같은 몇 가지 옵션이 있습니다.

옵션 1: TensorFlow Python에서 어댑터 코드를 작성하여 RNN 인터페이스를 Keras RNN 인터페이스에 맞춥니다. 즉, Keras LSTM 레이어에서 생성된 것과 동일한, 생성된 RNN 인터페이스의 함수에 대한 tf_implements 주석이 있는 tf.function입니다. 그런 다음 Keras LSTM에 사용된 것과 동일한 변환 API가 작동합니다.

옵션 2: 위의 방법을 사용할 수 없는 경우 (예: Keras LSTM에 현재 레이어 정규화와 같은 TensorFlow Lite의 통합된 LSTM 작업에 의해 노출된 일부 기능이 누락됨) 커스텀 변환 코드를 작성하여 TensorFlow Lite 변환기를 확장하고 여기에서 준비된 복합 함수 MLIR-pass에 연결합니다. 함수의 인터페이스는 API 계약처럼 취급해야 하며 통합 TensorFlow Lite LSTM 작업(예: 입력, 편향, 가중치, 투영, 레이어 정규화 등)으로 변환하는 데 필요한 인수를 포함해야 합니다. 이 함수에 인수로 전달된 텐서에 알려진 순위(예: MLIR의 RankedTensorType)가 권장됩니다. 이렇게 하면 이러한 텐서를 RankedTensorType으로 가정할 수 있는 변환 코드를 훨씬 쉽게 작성할 수 있으며, 통합 TensorFlow Lite 연산자의 피연산자에 상응하는 순위가 지정된 텐서로 변환하는 데 도움이 됩니다.

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

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

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

  2. 준비 복합 함수 패스를 확장하여 커스텀 복합 작업을 TensorFlow Lite 통합 LSTM 작업 변환에 연결합니다. LSTMCellSimple 변환 코드를 참고하세요.

    전환 계약:

  3. Weight와 projection 텐서가 전치됩니다.

  4. 전치된 가중치 텐서를 슬라이스하여 {input, recurrent}에서 {cell, input gate, ignore gate, output gate}를 추출합니다.

  5. {bias}부터 {cell, input gate, ignore gate, output gate}까지 편향 텐서를 분할하여 추출됩니다.

  6. projection은 전치된 프로젝션 텐서를 슬라이싱하여 추출됩니다.

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

  8. 정의된 모든 MLIR 패스와 TensorFlow Lite flatbuffer로의 최종 내보내기를 포함한 TensorFlow Lite 변환 인프라의 나머지 부분을 재사용할 수 있습니다.

알려진 문제/제한사항

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