Пользовательские операторы

Поскольку встроенная библиотека операторов TensorFlow Lite поддерживает только ограниченное количество операторов TensorFlow, не каждая модель может быть преобразована. Подробнее см. в разделе Совместимость с операторами .

Чтобы разрешить преобразование, пользователи могут предоставить собственную реализацию неподдерживаемого оператора TensorFlow в TensorFlow Lite, известную как пользовательский оператор. Если вместо этого вы хотите объединить ряд неподдерживаемых (или поддерживаемых) операторов TensorFlow в один объединенный оптимизированный пользовательский оператор, обратитесь к разделу объединения операторов .

Использование пользовательских операторов состоит из четырех шагов.

Давайте рассмотрим сквозной пример запуска модели с пользовательским оператором tf.atan (названным как Atan , см. #create_a_tensorflow_model), который поддерживается в TensorFlow, но не поддерживается в TensorFlow Lite.

Оператор TensorFlow Text является примером пользовательского оператора. Пример кода см. в учебнике «Преобразование TF Text в TF Lite» .

Пример: Пользовательский оператор Atan

Давайте рассмотрим пример поддержки оператора TensorFlow, которого нет в TensorFlow Lite. Предположим, что мы используем оператор Atan и строим очень простую модель для функции y = atan(x + offset) , где offset можно обучать.

Создайте модель TensorFlow

Следующий фрагмент кода обучает простую модель TensorFlow. Эта модель просто содержит пользовательский оператор с именем Atan , который представляет собой функцию y = atan(x + offset) , где offset можно обучать.

import tensorflow as tf

# Define training dataset and variables
x = [-8, 0.5, 2, 2.2, 201]
y = [-1.4288993, 0.98279375, 1.2490457, 1.2679114, 1.5658458]
offset = tf.Variable(0.0)

# Define a simple model which just contains a custom operator named `Atan`
@tf.function(input_signature=[tf.TensorSpec.from_tensor(tf.constant(x))])
def atan(x):
  return tf.atan(x + offset, name="Atan")

# Train model
optimizer = tf.optimizers.Adam(0.01)
def train(x, y):
    with tf.GradientTape() as t:
      predicted_y = atan(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: 0.99999905

На этом этапе, если вы попытаетесь создать модель TensorFlow Lite с флагами преобразователя по умолчанию, вы получите следующее сообщение об ошибке:

Error:
error: 'tf.Atan' op is neither a custom op nor a flex op.

Преобразование в модель TensorFlow Lite

Создайте модель TensorFlow Lite с пользовательскими операторами, установив атрибут конвертера allow_custom_ops , как показано ниже:

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

На этом этапе, если вы запустите его с интерпретатором по умолчанию, используя такие команды:

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

Вы все равно получите сообщение об ошибке:

Encountered unresolved custom op: Atan.

Создайте и зарегистрируйте оператора.

Все операторы TensorFlow Lite (как пользовательские, так и встроенные) определяются с помощью простого интерфейса на чистом 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;

Обратитесь к common.h для получения подробной информации о TfLiteContext и TfLiteNode . Первый предоставляет средства сообщения об ошибках и доступ к глобальным объектам, включая все тензоры. Последний позволяет реализациям получать доступ к своим входам и выходам.

Когда интерпретатор загружает модель, он вызывает init() один раз для каждого узла графа. Данный init() будет вызываться более одного раза, если операция используется в графе несколько раз. Для пользовательских операций будет предоставлен буфер конфигурации, содержащий flexbuffer, который сопоставляет имена параметров с их значениями. Буфер пуст для встроенных операций, потому что интерпретатор уже проанализировал параметры операции. Реализации ядра, которым требуется состояние, должны инициализировать его здесь и передать право собственности вызывающему объекту. Для каждого вызова 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, — это определить две функции ( Prepare и Eval ) и создать TfLiteRegistration :

TfLiteStatus AtanPrepare(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 AtanEval(TfLiteContext* context, TfLiteNode* node) {
  using namespace tflite;
  const TfLiteTensor* input = GetInput(context, node, 0);
  TfLiteTensor* output = GetOutput(context, node, 0);

  float* input_data = GetTensorData<float>(input);
  float* output_data = GetTensorData<float>(output);

  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] = atan(input_data[i]);
  }
  return kTfLiteOk;
}

TfLiteRegistration* Register_ATAN() {
  static TfLiteRegistration r = {nullptr, nullptr, AtanPrepare, AtanEval};
  return &r;
}

При инициализации OpResolver добавьте пользовательскую операцию в преобразователь (пример см. ниже). Это зарегистрирует оператор в Tensorflow Lite, чтобы TensorFlow Lite мог использовать новую реализацию. Обратите внимание, что последние два аргумента в TfLiteRegistration соответствуют функциям AtanPrepare и AtanEval , которые вы определили для пользовательской операции. Если бы вы использовали функции AtanInit и AtanFree для инициализации переменных, используемых в операции, и для освобождения места соответственно, то они были бы добавлены к первым двум аргументам 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("Atan", Register_ATAN());

Если набор встроенных операций считается слишком большим, новый OpResolver может быть сгенерирован кодом на основе заданного подмножества операций, возможно, только тех, которые содержатся в данной модели. Это эквивалент выборочной регистрации TensorFlow (простая версия доступна в каталоге tools ).

Если вы хотите определить свои пользовательские операторы в Java, вам в настоящее время необходимо создать свой собственный пользовательский слой JNI и скомпилировать свой собственный AAR в этом коде jni . Точно так же, если вы хотите определить эти операторы, доступные в Python, вы можете поместить свои регистрации в код-оболочку Python .

Обратите внимание, что процесс, аналогичный описанному выше, можно использовать для поддержки набора операций вместо одного оператора. Просто добавьте столько операторов AddCustom , сколько вам нужно. Кроме того, BuiltinOpResolver также позволяет переопределять реализации встроенных функций с помощью метода AddBuiltin .

Протестируйте и профилируйте своего оператора

Чтобы профилировать свою операцию с помощью инструмента тестирования TensorFlow Lite, вы можете использовать инструмент эталонной модели для TensorFlow Lite. В целях тестирования вы можете сделать так, чтобы ваша локальная сборка TensorFlow Lite знала о вашей пользовательской операции, добавив соответствующий вызов AddCustom (как показано выше) в register.cc .

Лучшие практики

  1. Осторожно оптимизируйте выделение и освобождение памяти. Выделение памяти в Prepare более эффективно, чем в Invoke , а выделение памяти перед циклом лучше, чем в каждой итерации. Используйте временные данные тензоров, а не настраивайте их самостоятельно (см. пункт 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 , т. е. эти макросы следует использовать до того, как будут выделены какие-либо ресурсы, которые могут привести к утечке.