Поскольку встроенная библиотека операторов LiteRT поддерживает лишь ограниченное количество операторов TensorFlow, не все модели можно преобразовать. Подробнее см. в разделе о совместимости операторов .
Чтобы обеспечить преобразование, пользователи могут предоставить собственную реализацию неподдерживаемого оператора TensorFlow в LiteRT, называемую пользовательским оператором. Если вместо этого вы хотите объединить ряд неподдерживаемых (или поддерживаемых) операторов TensorFlow в один объединённый оптимизированный пользовательский оператор, см. раздел «Объединение операторов» .
Использование пользовательских операторов состоит из четырех шагов.
Создайте модель TensorFlow. Убедитесь, что сохранённая модель (или определение графа) ссылается на правильно названный оператор LiteRT.
Конвертация в модель LiteRT. Для успешного преобразования модели убедитесь, что вы правильно установили атрибут конвертера LiteRT.
Создайте и зарегистрируйте оператор. Это необходимо для того, чтобы среда выполнения LiteRT знала, как сопоставлять ваш оператор и параметры в графе с исполняемым кодом C/C++.
Протестируйте и профилируйте свой оператор. Если вы хотите протестировать только свой собственный оператор, лучше всего создать модель только с вашим собственным оператором и использовать программу benchmark_model .
Давайте рассмотрим пример запуска модели с помощью пользовательского оператора tf.atan (названного Atan , см. раздел Создание модели TensorFlow ), который поддерживается в TensorFlow, но не поддерживается в LiteRT.
Оператор TensorFlow Text — пример пользовательского оператора. Пример кода см. в руководстве « Преобразование текста TF в LiteRT» .
Пример: пользовательский оператор Atan
Давайте рассмотрим пример поддержки оператора TensorFlow, которого нет в LiteRT. Предположим, мы используем оператор 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
На этом этапе, если вы попытаетесь сгенерировать модель LiteRT с флагами преобразователя по умолчанию, вы получите следующее сообщение об ошибке:
Error:
error: 'tf.Atan' op is neither a custom op nor a flex op.
Преобразовать в модель LiteRT
Создайте модель LiteRT с пользовательскими операторами, установив атрибут конвертера 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.
Создайте и зарегистрируйте оператора.
#include "third_party/tensorflow/lite/c/c_api.h"
#include "third_party/tensorflow/lite/c/c_api_opaque.h"
Пользовательские операторы LiteRT определяются с помощью простого API на чистом C, который состоит из непрозрачного типа ( TfLiteOperator ) и связанных функций.
TfLiteOperator — непрозрачный тип:
typedef struct TfLiteOperator TfLiteOperator;
TfLiteOperator хранит идентификатор и реализацию оператора. (Обратите внимание, что оператор отличается от своих операндов, которые хранятся в узлах графа LiteRT для узлов, вызывающих оператор.)
Экземпляры этого типа создаются с помощью вызовов TfLiteOperatorCreate и могут быть уничтожены с помощью вызова TfLiteOperatorDelete .
Идентификация оператора задается через параметры функции-конструктора TfLiteOperatorCreate :
TfLiteOperator*
TfLiteOperatorCreate(
TfLiteBuiltinOperator builtin_code, // Normally `TfLiteBuiltinCustom`.
const char* custom_name, // The name of the custom op.
int version // Normally `1` for the first version of a custom op.
);
Реализация оператора может определять «методы» со следующими сигнатурами. Все эти методы необязательны, но для успешного выполнения оператора реализация оператора должна определить и задать (с помощью функций-сеттеров) как минимум методы Prepare и Invoke .
// Initializes the op from serialized data.
void* Init(TfLiteOpaqueContext* context, const char* buffer, size_t length);
// Deallocates the op.
// The pointer `buffer` is the data previously returned by an Init invocation.
void Free(TfLiteOpaqueContext* context, void* buffer);
// Called when the inputs that this node depends on have been resized.
TfLiteStatus Prepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node);
// Called when the node is executed. (Should read node inputs and write to
// node outputs).
TfLiteStatus Invoke(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node);
// Retrieves the async kernel.
TfLiteAsyncKernel AsyncKernel(TfLiteOpaqueContext* context,
TfLiteOpaqueNode* node);
Имена функций (или префиксы пространств имён для C++) в вашей реализации операции не обязательно должны совпадать с именами функций в приведённом выше фрагменте кода, поскольку API пользовательских операций TF Lite будет использовать только их адреса. Более того, мы рекомендуем объявлять их в анонимном пространстве имён или как статические функции.
Однако хорошей идеей будет включить имя вашего оператора в качестве пространства имен или префикса в имена этих функций:
С++
namespace my_namespace::my_custom_op { void* Init(TfLiteOpaqueContext* context, const char* buffer, size_t length) { ... } // ... plus definitions of Free, Prepare, and Invoke ... }
С
void* MyCustomOpInit(TfLiteOpaqueContext* context, const char* buffer, size_t length) { ... } // ... plus definitions of MyCustomOpFree, MyCustomOpPrepare, and // MyCustomOpInvoke.
Поскольку это API C, эти «методы» реализованы как указатели на функции C в типе TfLiteOperator , которые задаются путем передачи адресов ваших функций реализации соответствующим функциям-сеттерам TfLiteOperatorSet MethodName :
void TfLiteOperatorSetInit(
TfLiteOperator* operator,
void* (*init)(TfLiteOpaqueContext* context, const char* buffer,
size_t length));
void TfLiteOperatorSetFree(
TfLiteOperator* operator,
void (*free)(TfLiteOpaqueContext* context, void* data));
void TfLiteOperatorSetPrepare(
TfLiteOperator* operator,
TfLiteStatus (*prepare)(TfLiteOpaqueContext* context,
TfLiteOpaqueNode* node));
void TfLiteOperatorSetInvoke(
TfLiteOperator* operator,
TfLiteStatus (*invoke)(TfLiteOpaqueContext* context,
TfLiteOpaqueNode* node));
void TfLiteOperatorSetAsyncKernel(
TfLiteOperator* operator,
struct TfLiteAsyncKernel* (*async_kernel)(TfLiteOpaqueContext* context,
TfLiteOpaqueNode* node));
Подробную информацию о TfLiteContext и TfLiteNode см. в common.h . TfLiteContext предоставляет возможности сообщения об ошибках и доступ к глобальным объектам, включая все тензоры. TfLiteNode позволяет реализациям операторов получать доступ к своим входным и выходным данным.
Когда интерпретатор загружает модель, он вызывает метод Init() один раз для каждого узла графа. Заданный метод Init() будет вызван более одного раза, если операция используется в графе несколько раз. Для пользовательских операций будет предоставлен буфер конфигурации, содержащий гибкий буфер, который сопоставляет имена параметров с их значениями. Для встроенных операций этот буфер пуст, поскольку интерпретатор уже проанализировал параметры операции. Реализации ядра, которым требуется состояние, должны инициализировать его здесь и передать владение вызывающей стороне. Для каждого вызова Init() будет соответствующий вызов Free() , что позволит реализациям освободить буфер, который они могли бы выделить в Init() .
При каждом изменении размера входных тензоров интерпретатор проходит по графу, уведомляя реализации об этом изменении. Это даёт им возможность изменить размер внутреннего буфера, проверить корректность входных форм и типов и пересчитать выходные формы. Всё это выполняется с помощью метода Prepare() , а реализации могут получить доступ к их состоянию с помощью TfLiteOpaqueNodeGetUserData(node) .
Наконец, каждый раз при запуске вывода интерпретатор обходит граф, вызывая метод Invoke() , и здесь состояние также доступно как TfLiteOpaqueNodeGetUserData(node) .
Пользовательские операции можно реализовать, определив эти функции-«методы», а затем определив функцию, которая возвращает экземпляр TfLiteOperator , созданный путем вызова TfLiteOperatorCreate , а затем соответствующих методов установки:
С++
namespace my_namespace::my_custom_op { namespace { void* Init(TfLiteOpaqueContext* context, const char* buffer, size_t length) { ... } void Free(TfLiteOpaqueContext* context, void* buffer) { ... } TfLiteStatus Prepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { ... } TfLiteStatus Invoke(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {... } }; const TfLiteOperator* MyCustomOperator() { // Singleton instance, intentionally never destroyed. static const TfLiteOperator* my_custom_op = ()[] { TfLiteOperator* r = TfLiteOperatorCreate( kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1); TfLiteOperatorSetInit(r, Init); TfLiteOperatorSetFree(r, Free); TfLiteOperatorSetPrepare(r, Prepare); TfLiteOperatorSetInvoke(r, Eval); return r; }; return my_custom_op; } } // namespace my_namespace
С
static void* MyCustomOpInit(TfLiteOpaqueContext* context, const char* buffer, size_t length) { ... } static void MyCustomOpFree(TfLiteOpaqueContext* context, void* buffer) { ... } static TfLiteStatus MyCustomOpPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { ... } static TfLiteStatus MyCustomOpInvoke(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) {... } static TfLiteOperator* MyCustomOpCreate() { const TfLiteOperator* r = TfLiteOperatorCreate( kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1); TfLiteOperatorSetInit(r, MyCustomOpInit); TfLiteOperatorSetFree(r, MyCustomOpFree); TfLiteOperatorSetPrepare(r, MyCustomOpPrepare); TfLiteOperatorSetInvoke(r, MyCustomOpEval); return r; } const TfLiteOperator* MyCustomOperator() { // Singleton instance, intentionally never destroyed. static const TfLiteOperator* my_custom_op = MyCustomOpCreate(); return my_custom_op; }
Обратите внимание, что регистрация не происходит автоматически, и необходимо явно вызвать функцию MyCustomOperator (подробности см. ниже). В то время как стандартный BuiltinOpResolver (доступный из цели :builtin_ops ) отвечает за регистрацию встроенных операций, пользовательские операции придётся собирать в отдельных библиотеках.
Определение ядра в среде выполнения LiteRT
Все, что нам нужно сделать для использования OP в LiteRT, — это определить две функции ( Prepare и Eval ) и третью для построения TfLiteOperator :
С++
namespace atan_op { namespace { TfLiteStatus AtanPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumInputs(node), 1); TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumOutputs(node), 1); const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0); TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0); int num_dims = TfLiteOpaqueTensorNumDimensions(input); TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims); for (int i=0; i < num_dims; ++i) { output_size->data[i] = input->dims->data[i]; } return TfLiteOpaqueContextResizeTensor(context, output, output_size); } TfLiteStatus AtanEval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0); TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0); float* input_data = static_cast<float*>(TfLiteOpaqueTensorData(input)); float* output_data = static_cast<float*>(TfLiteOpaqueTensorData(output)); size_t count = 1; int num_dims = TfLiteOpaqueTensorNumDimensions(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; } } // anonymous namespace const TfLiteOperator* AtanOperator() { // Singleton instance, intentionally never destroyed. static const TfLiteOperator* atan_op = ()[] { auto* r = TfLiteOperatorCreate( kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1); TfLiteOperatorSetPrepare(r, Prepare); TfLiteOperatorSetInvoke(r, Eval); return r; }; return atan_op; } } // namespace atan_op
С
static TfLiteStatus AtanPrepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumInputs(node), 1); TF_LITE_OPAQUE_ENSURE_EQ(context, TfLiteOpaqueNodeNumOutputs(node), 1); const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0); TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0); int num_dims = TfLiteOpaqueTensorNumDimensions(input); TfLiteIntArray* output_size = TfLiteIntArrayCreate(num_dims); for (int i = 0; i < num_dims; ++i) { output_size->data[i] = input->dims->data[i]; } return TfLiteOpaqueContextResizeTensor(context, output, output_size); } static TfLiteStatus AtanEval(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { const TfLiteOpaqueTensor* input = TfLiteOpaqueNodeGetInput(context, node, 0); TfLiteOpaqueTensor* output = TfLiteOpaqueNodeGetOutput(context, node, 0); float* input_data = static_cast<float*>(TfLiteOpaqueTensorData(input)); float* output_data = static_cast<float*>(TfLiteOpaqueTensorData(output)); size_t count = 1; int num_dims = TfLiteOpaqueTensorNumDimensions(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; } static const TfLiteOperator* AtanOpCreate() { TfLiteOperator* r = TfLiteOperatorCreate( kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1); TfLiteOperatorSetPrepare(r, Prepare); TfLiteOperatorSetInvoke(r, Eval); return r; } const TfLiteOperator* AtanOperator() { // Singleton instance, intentionally never destroyed. static const TfLiteOperator* atan_op = AtanOpCreate(); return atan_op; }
При инициализации OpResolver добавьте пользовательскую операцию в резолвер (см. пример ниже). Это зарегистрирует оператор в LiteRT, чтобы LiteRT мог использовать новую реализацию.
Зарегистрировать оператор в библиотеке ядра
Теперь нам нужно зарегистрировать оператор в библиотеке ядра. Это делается с помощью OpResolver . Интерпретатор автоматически загрузит библиотеку ядер, которая будет назначена для выполнения каждого из операторов в модели. Хотя библиотека по умолчанию содержит только встроенные ядра, её можно заменить/дополнить пользовательской библиотекой операторов.
Класс OpResolver , который преобразует коды и имена операторов в фактический код, определяется следующим образом:
class OpResolver {
public:
virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
virtual TfLiteRegistration* FindOp(const char* op) const = 0;
...
};
Обратите внимание, что для обратной совместимости этот класс использует старый конкретный тип TfLiteRegistration вместо непрозрачного типа TfLiteOperator , но структура TfLiteRegistration содержит поле registration_external типа TfLiteOperator* .
Классы MutableOpResolver и BuiltinOpResolver являются производными от OpResolver :
class MutableOpResolver : public OpResolver {
public:
MutableOpResolver(); // Constructs an initially empty op resolver.
void AddAll(const MutableOpResolver& other);
...
};
class BuiltinOpResolver : public MutableOpResolver {
public:
BuiltinOpResolver(); // Constructs an op resolver with all the builtin ops.
};
Для обычного использования (без пользовательских операций) необходимо использовать BuiltinOpResolver и написать:
tflite::ops::builtin::BuiltinOpResolver resolver;
Чтобы добавить созданную выше пользовательскую операцию, вы можете вместо этого использовать MutableOpResolver и вызвать tflite::AddOp (до передачи решателя в InterpreterBuilder ):
tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
tflite::AddOp(&resolver, AtanOpRegistration());
Если набор встроенных операций окажется слишком большим, можно сгенерировать новый OpResolver на основе заданного подмножества операций, возможно, только тех, которые содержатся в данной модели. Это эквивалентно выборочной регистрации TensorFlow (и её простая версия доступна в каталоге tools ).
Если вы хотите определить собственные операторы в Java, вам в настоящее время потребуется создать собственный слой JNI и скомпилировать свой AAR в этом коде JNI . Аналогично, если вы хотите определить эти операторы, доступные в Python, вы можете разместить свои регистрации в коде-обёртке Python .
Обратите внимание, что аналогичный описанному выше процесс можно использовать для поддержки набора операций вместо одного оператора. Просто добавьте столько операторов AddCustom , сколько необходимо. Кроме того, MutableOpResolver позволяет переопределять реализации встроенных функций с помощью AddBuiltin .
Протестируйте и составьте профиль своего оператора
Для профилирования вашей операции с помощью инструмента бенчмарка LiteRT вы можете использовать инструмент бенчмарк-моделирования для LiteRT. Для тестирования вы можете настроить локальную сборку LiteRT на использование вашей пользовательской операции, добавив соответствующий вызов AddCustom (как показано выше) в файл register.cc.
Лучшие практики
Тщательно оптимизируйте выделение и освобождение памяти. Выделение памяти в
Prepareэффективнее, чем вInvoke, а выделение памяти перед циклом — лучше, чем в каждой итерации. Используйте временные данные тензоров вместо самостоятельного выделения памяти (см. пункт 2). Используйте указатели/ссылки вместо копирования как можно большего объёма памяти.Если структура данных будет сохраняться в течение всей операции, рекомендуется предварительно выделить память с помощью временных тензоров. Для ссылки на индексы тензоров в других функциях может потребоваться структура OpData. См. пример в ядре для свёртки . Пример кода приведён ниже.
struct MyOpData { int temp_tensor_index; ... }; void* Init(TfLiteOpaqueContext* context, const char* buffer, size_t length) { auto* op_data = new MyOpData{}; ... return op_data; } void Free(TfLiteOpaqueContext* context, void* buffer) { ... delete reinterpret_cast<MyOpData*>(buffer); } TfLiteStatus Prepare(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { ... auto* op_data = reinterpret_cast<MyOpData*>(TfLiteOpaqueNodeGetUserData(node)); const int num_temporaries = 1; int temporary_tensor_indices[num_temporaries]; TfLiteOpaqueTensorBuilder* builder = TfLiteOpaqueTensorBuilderCreate(); TfLiteOpaqueTensorBuilderSetType(builder, kTfLiteFloat32); TfLiteOpaqueTensorBuilderSetAllocationType(builder, kTfLiteArenaRw); TfLiteOpaqueContextAddTensor(context, builder, &temporary_tensor_indices[0]); TfLiteOpaqueTensorBuilderDelete(builder); TfLiteOpaqueNodeSetTemporaries(node, temporary_tensor_indices, num_temporaries); op_data->temp_tensor_index = temporary_tensor_indices[0]; ... return kTfLiteOk; } TfLiteStatus Invoke(TfLiteOpaqueContext* context, TfLiteOpaqueNode* node) { ... auto* op_data = reinterpret_cast<MyOpData*>( TfLiteOpaqueNodeGetUserData(node)); TfLiteOpaqueTensor* temp_tensor = TfLiteOpaqueContextGetOpaqueTensor(context, op_data->temp_tensor_index); TF_LITE_OPAQUE_ENSURE(context, TfLiteTensorType(temp_tensor) == kTfLiteFloat32); TF_LITE_OPAQUE_ENSURE(context, TfLiteTensorGetAllocationType(temp_Tensor) == kTfLiteArenaRw); void *temp_data = TfLiteTensorData(temp_tensor); TF_LITE_OPAQUE_ENSURE(context, temp_data != nullptr); ... return kTfLiteOk; }Если это не приведет к слишком большим затратам памяти, предпочтительнее использовать статический массив фиксированного размера (или предварительно выделенный
std::vectorвResize) вместо использования динамически выделяемогоstd::vectorна каждой итерации выполнения.Избегайте создания экземпляров шаблонов контейнеров стандартной библиотеки, которых ещё нет, поскольку они влияют на размер двоичного кода. Например, если в вашей операции требуется
std::map, которого нет в других ядрах, использованиеstd::vectorс прямым индексированием может быть эффективным, сохраняя при этом небольшой размер двоичного кода. Посмотрите, как используются другие ядра, чтобы получить более подробную информацию (или спросите).Проверьте указатель на память, возвращаемый функцией
malloc. Если этот указатель равенnullptr, с ним не следует выполнять никаких операций. Если при выполненииmallocв функции возникает ошибка выхода, перед выходом освободите память.Используйте
TF_LITE_OPAQUE_ENSURE(context, condition)для проверки определённого условия. Ваш код не должен оставлять память зависшей при использованииTF_LITE_OPAQUE_ENSURE, то есть эти макросы следует использовать до выделения любых ресурсов, которые могут привести к утечке.