Google I/O는 끝입니다! TensorFlow 세션 확인하기 세션 보기

사용자 정의 연산자

TensorFlow Lite 내장 연산자 라이브러리는 제한된 수의 TensorFlow 연산자만 지원하므로 모든 모델을 변환할 수 있는 것은 아닙니다. 자세한 내용은 운영자 호환성 을 참조하십시오.

변환을 허용하기 위해 사용자는 사용자 지정 연산자로 알려진 TensorFlow Lite에서 지원되지 않는 TensorFlow 연산자의 사용자 지정 구현을 제공할 수 있습니다. 대신 지원되지 않는(또는 지원되는) 일련의 TensorFlow 연산자를 통합된 최적화된 단일 사용자 지정 연산자로 결합하려는 경우 연산자 융합 을 참조하세요.

사용자 지정 연산자 사용은 4단계로 구성됩니다.

TensorFlow에서는 지원되지만 TensorFlow Lite에서는 지원되지 않는 사용자 지정 연산자 tf.sin ( Sin 으로 명명, #create_a_tensorflow_model 참조)을 사용하여 모델을 실행하는 종단 간 예제를 살펴보겠습니다.

예: 사용자 정의 Sin 연산자

TensorFlow Lite에는 없는 TensorFlow 연산자를 지원하는 예를 살펴보겠습니다. Sin 연산자를 사용하고 있고 offset 이 훈련 가능한 함수 y = sin(x + offset) 에 대한 매우 간단한 모델을 구축하고 있다고 가정합니다.

TensorFlow 모델 만들기

다음 코드 스니펫은 간단한 TensorFlow 모델을 학습시킵니다. 이 모델에는 함수 y = sin(x + offset)Sin 이라는 사용자 지정 연산자가 포함되어 있습니다. 여기서 offset 은 훈련 가능합니다.

import tensorflow as tf

# Define training dataset and variables
x = [-8, 0.5, 2, 2.2, 201]
y = [-0.6569866 ,  0.99749499,  0.14112001, -0.05837414,  0.80641841]
offset = tf.Variable(0.0)

# Define a simple model which just contains a custom operator named `Sin`
@tf.function
def sin(x):
  return tf.sin(x + offset, name="Sin")

# Train model
optimizer = tf.optimizers.Adam(0.01)
def train(x, y):
    with tf.GradientTape() as t:
      predicted_y = sin(x)
      loss = tf.reduce_sum(tf.square(predicted_y - y))
    grads = t.gradient(loss, [offset])
    optimizer.apply_gradients(zip(grads, [offset]))

for i in range(1000):
    train(x, y)

print("The actual offset is: 1.0")
print("The predicted offset is:", offset.numpy())
The actual offset is: 1.0
The predicted offset is: 1.0000001

이 시점에서 기본 변환기 플래그를 사용하여 TensorFlow Lite 모델을 생성하려고 하면 다음 오류 메시지가 표시됩니다.

Error:
Some of the operators in the model are not supported by the standard TensorFlow
Lite runtime...... Here is
a list of operators for which you will need custom implementations: Sin.

TensorFlow Lite 모델로 변환

아래와 같이 변환기 속성 allow_custom_ops 를 설정하여 사용자 지정 연산자가 있는 TensorFlow Lite 모델을 만듭니다.

converter = tf.lite.TFLiteConverter.from_concrete_functions([sin.get_concrete_function(x)], sin)
converter.allow_custom_ops = True
tflite_model = converter.convert()

이때 기본 인터프리터로 실행하면 다음과 같은 오류 메시지가 표시됩니다.

Error:
Didn't find custom operator for name 'Sin'
Registration failed.

연산자를 만들고 등록합니다.

모든 TensorFlow Lite 연산자(맞춤형 및 내장형 모두)는 4가지 기능으로 구성된 간단한 순수 C 인터페이스를 사용하여 정의됩니다.

typedef struct {
  void* (*init)(TfLiteContext* context, const char* buffer, size_t length);
  void (*free)(TfLiteContext* context, void* buffer);
  TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node);
  TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node);
} TfLiteRegistration;

TfLiteContextTfLiteNode 에 대한 자세한 내용은 common.h 를 참조하십시오. 전자는 오류 보고 기능과 모든 텐서를 포함한 전역 개체에 대한 액세스를 제공합니다. 후자는 구현이 입력 및 출력에 액세스할 수 있도록 합니다.

인터프리터는 모델을 로드할 때 그래프의 각 노드에 대해 한 번씩 init() 를 호출합니다. 주어진 init() 는 그래프에서 op가 여러 번 사용되는 경우 두 번 이상 호출됩니다. 사용자 지정 작업의 경우 매개변수 이름을 해당 값에 매핑하는 플렉스 버퍼를 포함하는 구성 버퍼가 제공됩니다. 인터프리터가 이미 연산 매개변수를 구문 분석했기 때문에 버퍼는 내장 연산에 대해 비어 있습니다. 상태가 필요한 커널 구현은 여기에서 초기화하고 소유권을 호출자에게 이전해야 합니다. 각 init() 호출에 대해 free() 에 대한 해당 호출이 있으므로 구현에서 init() 에 할당했을 수 있는 버퍼를 삭제할 수 있습니다.

입력 텐서의 크기가 조정될 때마다 인터프리터는 변경 구현을 알리는 그래프를 살펴봅니다. 이를 통해 내부 버퍼의 크기를 조정하고, 입력 모양 및 유형의 유효성을 확인하고, 출력 모양을 다시 계산할 수 있습니다. 이것은 모두 prepare() 를 통해 수행되며 구현은 node->user_data 를 사용하여 상태에 액세스할 수 있습니다.

마지막으로, 추론이 실행될 때마다 인터프리터는 invoke() 를 호출하는 그래프를 횡단하며 여기에서도 상태를 node->user_data 로 사용할 수 있습니다.

사용자 지정 작업은 내장 작업과 정확히 동일한 방식으로 구현될 수 있습니다. 이러한 네 가지 기능과 일반적으로 다음과 같은 전역 등록 기능을 정의하면 됩니다.

namespace tflite {
namespace ops {
namespace custom {
  TfLiteRegistration* Register_MY_CUSTOM_OP() {
    static TfLiteRegistration r = {my_custom_op::Init,
                                   my_custom_op::Free,
                                   my_custom_op::Prepare,
                                   my_custom_op::Eval};
    return &r;
  }
}  // namespace custom
}  // namespace ops
}  // namespace tflite

등록은 자동이 아니며 Register_MY_CUSTOM_OP 를 명시적으로 호출해야 합니다. 표준 BuiltinOpResolver ( :builtin_ops 대상에서 사용 가능)가 내장 등록을 처리하는 동안 사용자 지정 작업은 별도의 사용자 지정 라이브러리에서 수집해야 합니다.

TensorFlow Lite 런타임에서 커널 정의

TensorFlow Lite에서 op를 사용하기 위해 해야 할 일은 두 가지 함수( PrepareEval )를 정의하고 TfLiteRegistration 을 구성하는 것입니다.

TfLiteStatus SinPrepare(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  const TfLiteTensor* input = GetInput(context, node, 0);
  TfLiteTensor* output = GetOutput(context, node, 0);

  int num_dims = NumDimensions(input);

  TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims);
  for (int i=0; i<num_dims; ++i) {
    output_size->data[i] = input->dims->data[i];
  }

  return context->ResizeTensor(context, output, output_size);
}

TfLiteStatus SinEval(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  const TfLiteTensor* input = GetInput(context, node,0);
  TfLiteTensor* output = GetOutput(context, node,0);

  float* input_data = input->data.f;
  float* output_data = output->data.f;

  size_t count = 1;
  int num_dims = NumDimensions(input);
  for (int i = 0; i < num_dims; ++i) {
    count *= input->dims->data[i];
  }

  for (size_t i=0; i<count; ++i) {
    output_data[i] = sin(input_data[i]);
  }
  return kTfLiteOk;
}

TfLiteRegistration* Register_SIN() {
  static TfLiteRegistration r = {nullptr, nullptr, SinPrepare, SinEval};
  return &r;
}

OpResolver 를 초기화할 때 사용자 지정 작업을 확인자에 추가합니다(아래 예제 참조). 그러면 TensorFlow Lite가 새 구현을 사용할 수 있도록 연산자가 Tensorflow Lite에 등록됩니다. TfLiteRegistration 의 마지막 두 인수는 사용자 지정 작업에 대해 정의한 SinPrepareSinEval 함수에 해당합니다. SinInitSinFree 함수를 사용하여 각각 op에서 사용되는 변수를 초기화하고 공간을 확보한 경우 TfLiteRegistration 의 처음 두 인수에 추가됩니다. 이러한 인수는 이 예에서 nullptr 로 설정됩니다.

커널 라이브러리에 연산자 등록

이제 커널 라이브러리에 연산자를 등록해야 합니다. 이것은 OpResolver 로 수행됩니다. 배후에서 인터프리터는 모델의 각 연산자를 실행하기 위해 할당될 커널 라이브러리를 로드합니다. 기본 라이브러리에는 내장 커널만 포함되어 있지만 사용자 지정 라이브러리 연산 연산자로 교체/증가할 수 있습니다.

연산자 코드와 이름을 실제 코드로 변환하는 OpResolver 클래스는 다음과 같이 정의됩니다.

class OpResolver {
  virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
  virtual TfLiteRegistration* FindOp(const char* op) const = 0;
  virtual void AddBuiltin(tflite::BuiltinOperator op, TfLiteRegistration* registration) = 0;
  virtual void AddCustom(const char* op, TfLiteRegistration* registration) = 0;
};

정기적으로 사용하려면 BuiltinOpResolver 를 사용하고 다음을 작성해야 합니다.

tflite::ops::builtin::BuiltinOpResolver resolver;

위에서 생성한 사용자 지정 작업을 추가하려면 AddOp 를 호출합니다(해석기를 InterpreterBuilder 에 전달하기 전에).

resolver.AddCustom("Sin", Register_SIN());

내장 작업 세트가 너무 큰 것으로 간주되면 주어진 모델에 포함된 작업만 가능한 작업의 주어진 하위 집합을 기반으로 새 OpResolver 를 코드 생성할 수 있습니다. 이것은 TensorFlow의 선택적 등록과 동일합니다(단순한 버전은 tools 디렉토리에서 사용할 수 있습니다).

Java에서 사용자 정의 연산자를 정의하려면 현재 사용자 정의 JNI 계층을 빌드 하고 이 jni 코드에서 자체 AAR을 컴파일해야 합니다. 마찬가지로 Python에서 사용할 수 있는 이러한 연산자를 정의하려면 Python 래퍼 코드 에 등록을 배치할 수 있습니다.

단일 연산자 대신 일련의 작업을 지원하기 위해 위와 유사한 프로세스를 따를 수 있습니다. 필요한 만큼 AddCustom 연산자를 추가하기만 하면 됩니다. 또한 BuiltinOpResolver 를 사용하면 AddBuiltin 을 사용하여 내장 구현을 재정의할 수도 있습니다.

운영자 테스트 및 프로파일링

TensorFlow Lite 벤치마크 도구로 작업을 프로파일링하려면 TensorFlow Lite용 벤치 마크 모델 도구 를 사용할 수 있습니다. 테스트 목적을 위해 register.cc 에 적절한 AddCustom 호출(위 참조)을 추가하여 TensorFlow Lite의 로컬 빌드가 사용자 정의 작업을 인식하도록 할 수 있습니다.

모범 사례

  1. 메모리 할당 및 할당 해제를 신중하게 최적화하십시오. Prepare 에서 메모리를 할당하는 것이 Invoke 에서보다 더 효율적이며 루프 전에 메모리를 할당하는 것이 모든 반복에서보다 좋습니다. 자신을 mallocing하는 대신 임시 텐서 데이터를 사용합니다(항목 2 참조). 가능한 한 많이 복사하는 대신 포인터/참조를 사용하십시오.

  2. 전체 작업 동안 데이터 구조가 지속되는 경우 임시 텐서를 사용하여 메모리를 미리 할당하는 것이 좋습니다. 다른 함수에서 텐서 인덱스를 참조하려면 OpData 구조체를 사용해야 할 수도 있습니다. 컨볼루션에 대한 커널 의 예를 참조하십시오. 샘플 코드 스니펫은 다음과 같습니다.

    auto* op_data = reinterpret_cast<OpData*>(node->user_data);
    TfLiteIntArrayFree(node->temporaries);
    node->temporaries = TfLiteIntArrayCreate(1);
    node->temporaries->data[0] = op_data->temp_tensor_index;
    TfLiteTensor* temp_tensor = &context->tensors[op_data->temp_tensor_index];
    temp_tensor->type =  kTfLiteFloat32;
    temp_tensor->allocation_type = kTfLiteArenaRw;
    
  3. 너무 많은 메모리 낭비가 발생하지 않는다면 실행 반복마다 동적으로 할당된 std::vector 를 사용하는 것보다 정적 고정 크기 배열(또는 Resize 에서 미리 할당된 std::vector )을 사용하는 것을 선호합니다.

  4. 바이너리 크기에 영향을 미치므로 이미 존재하지 않는 표준 라이브러리 컨테이너 템플릿을 인스턴스화하지 마십시오. 예를 들어, 다른 커널에는 없는 작업에 std::map 이 필요한 경우 직접 인덱싱 매핑과 함께 std::vector 를 사용하면 바이너리 크기를 작게 유지하면서 작동할 수 있습니다. 통찰력을 얻기 위해(또는 질문하기 위해) 다른 커널이 무엇을 사용하는지 확인하십시오.

  5. malloc 에 의해 반환된 메모리에 대한 포인터를 확인하십시오. 이 포인터가 nullptr 인 경우 해당 포인터를 사용하여 작업을 수행해서는 안 됩니다. 함수에서 malloc 을 수행하고 오류 종료가 있는 경우 종료하기 전에 메모리 할당을 해제하십시오.

  6. TF_LITE_ENSURE(context, condition) 를 사용하여 특정 조건을 확인합니다. TF_LITE_ENSURE 를 사용할 때 코드에서 메모리가 정지된 상태로 남아 있어서는 안 됩니다. 즉, 누수될 리소스가 할당되기 전에 이러한 매크로를 사용해야 합니다.