由於 TensorFlow Lite 內建運算子程式庫僅支援少數 TensorFlow 運算子,並非所有模型都可轉換。如需詳細資訊,請參閱運算子相容性。
若要允許轉換,使用者可以在 TensorFlow Lite 中提供不受支援 TensorFlow 運算子的自訂實作,稱為自訂運算子。如果您希望將一系列不受支援 (或受支援) 的 TensorFlow 運算子合併為單一融合最佳化的自訂運算子,請參閱運算子融合。
使用自訂運算子包含四個步驟。
建立 TensorFlow 模型。 請確認 Saved Model (或 Graph Def) 參照的是正確命名的 TensorFlow Lite 運算子。
轉換為 TensorFlow Lite 模型。 請確認您已設定正確的 TensorFlow Lite 轉換器屬性,以便順利轉換模型。
建立並註冊運算子。 這樣 TensorFlow Lite 執行階段才知道如何將圖形中的運算子和參數對應至可執行的 C/C++ 程式碼。
測試及分析您的運算子。 如果您只想測試自訂運算子,最好建立一個僅包含自訂運算子的模型,並使用 benchmark_model 程式。
讓我們逐步瞭解執行模型的端對端範例,該模型具有自訂運算子 tf.atan
(命名為 Atan
,請參閱 #create_a_tensorflow_model),該運算子在 TensorFlow 中受支援,但在 TensorFlow Lite 中不受支援。
TensorFlow Text 運算子是自訂運算子的範例。如需程式碼範例,請參閱將 TF Text 轉換為 TF Lite 教學課程。
範例:自訂 Atan
運算子
讓我們逐步瞭解支援 TensorFlow Lite 沒有的 TensorFlow 運算子的範例。假設我們使用 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.
建立並註冊運算子。
#include "tensorflow/lite/c/c_api.h"
#include "tensorflow/lite/c/c_api_opaque.h"
TensorFlow Lite 自訂運算子是使用簡單的純 C API 定義的,該 API 包含不透明類型 (TfLiteRegistrationExternal
) 和相關函數。
TfLiteRegistrationExternal
是不透明類型
typedef struct TfLiteRegistrationExternal TfLiteRegistrationExternal;
TfLiteRegistrationExternal
儲存運算子的身分和實作。(請注意,運算子與其運算元不同,運算元儲存在呼叫運算子的節點的 TF Lite 圖形節點中。)
此類型的執行個體是透過呼叫 TfLiteRegistrationExternalCreate
建構的,並且可以透過呼叫 TfLiteRegistrationExternalDelete
銷毀。
運算子的身分是透過建構函式 TfLiteRegistrationExternalCreate
的參數設定的
TfLiteRegistrationExternal*
TfLiteRegistrationExternalCreate(
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.
);
運算子實作可以定義具有以下簽名的「方法」。所有這些方法都是選用的,但為了成功評估運算子,運算子實作需要定義和設定(使用 setter 函數)至少 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++)不必與上述程式碼片段中的函數名稱相符,因為 TF Lite 自訂運算元 API 只會使用它們的位址。實際上,我們建議您在匿名命名空間中或作為靜態函數宣告它們。
但最好將您的運算子名稱包含在這些函數名稱的命名空間或前綴中
C++
namespace my_namespace::my_custom_op { void* Init(TfLiteOpaqueContext* context, const char* buffer, size_t length) { ... } // ... plus definitions of Free, Prepare, and Invoke ... }
C
void* MyCustomOpInit(TfLiteOpaqueContext* context, const char* buffer, size_t length) { ... } // ... plus definitions of MyCustomOpFree, MyCustomOpPrepare, and // MyCustomOpInvoke.
由於這是 C API,因此這些「方法」實作為 TfLiteRegistrationExternal
類型中的 C 函數指標,這些指標是透過將您的實作函數的位址傳遞給對應的 setter 函數 TfLiteRegistrationExternalSet
MethodName 來設定的
void TfLiteRegistrationExternalSetInit(
TfLiteRegistrationExternal* registration,
void* (*init)(TfLiteOpaqueContext* context, const char* buffer,
size_t length));
void TfLiteRegistrationExternalSetFree(
TfLiteRegistrationExternal* registration,
void (*free)(TfLiteOpaqueContext* context, void* data));
void TfLiteRegistrationExternalSetPrepare(
TfLiteRegistrationExternal* registration,
TfLiteStatus (*prepare)(TfLiteOpaqueContext* context,
TfLiteOpaqueNode* node));
void TfLiteRegistrationExternalSetInvoke(
TfLiteRegistrationExternal* registration,
TfLiteStatus (*invoke)(TfLiteOpaqueContext* context,
TfLiteOpaqueNode* node));
void TfLiteRegistrationExternalSetAsyncKernel(
TfLiteRegistrationExternal* registration,
struct TfLiteAsyncKernel* (*async_kernel)(TfLiteOpaqueContext* context,
TfLiteOpaqueNode* node));
如需 TfLiteContext
和 TfLiteNode
的詳細資訊,請參閱 common.h
。TfLiteContext
提供錯誤報告功能和對全域物件的存取權,包括所有張量。TfLiteNode
允許運算子實作存取其輸入和輸出。
當解譯器載入模型時,它會針對圖形中的每個節點呼叫 Init()
方法一次。如果運算元在圖形中多次使用,則給定的 Init()
將被呼叫多次。對於自訂運算元,將提供一個組態緩衝區,其中包含將參數名稱對應到其值的 flexbuffer。對於內建運算元,緩衝區為空,因為解譯器已剖析運算元參數。需要狀態的核心實作應在此處初始化它,並將所有權轉移給呼叫者。對於每個 Init()
呼叫,都會有一個對應的 Free()
呼叫,允許實作處置它們可能在 Init()
中配置的緩衝區。
每當輸入張量調整大小時,解譯器都會遍歷圖形,通知實作變更。這讓他們有機會調整其內部緩衝區大小、檢查輸入形狀和類型的有效性,並重新計算輸出形狀。所有這些都是透過 Prepare()
方法完成的,實作可以使用 TfLiteOpaqueNodeGetUserData(node)
存取其狀態。
最後,每次執行推論時,解譯器都會遍歷圖形,呼叫 Invoke()
方法,並且此處的狀態也以 TfLiteOpaqueNodeGetUserData(node)
的形式提供。
自訂運算元可以透過定義那些「方法」函數來實作,然後定義一個函數,該函數傳回透過呼叫 TfLiteRegistrationExternalCreate
然後相關的 setter 方法建構的 TfLiteRegistrationExternal
執行個體
C++
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 TfLiteRegistrationExternal* MyCustomOpRegistrationExternal() { // Singleton instance, intentionally never destroyed. static const TfLiteRegistrationExternal* my_custom_op = ()[] { TfLiteRegistrationExternal* r = TfLiteRegistrationExternalCreate( kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1); TfLiteRegistrationExternalSetInit(r, Init); TfLiteRegistrationExternalSetFree(r, Free); TfLiteRegistrationExternalSetPrepare(r, Prepare); TfLiteRegistrationExternalSetInvoke(r, Eval); return r; }; return my_custom_op; } const TfLiteRegistration* MyCustomOpRegistration() { static const TfLiteRegistration my_custom_op { .registration_external = MyCustomOpRegistrationExternal(); }; return my_custom_op; } } // namespace my_namespace
C
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 TfLiteRegistrationExternal* MyCustomOpCreate() { const TfLiteRegistrationExternal* r = TfLiteRegistrationExternalCreate( kTfLiteBuiltinCustom, "MyCustomOp", /*version=*/ 1); TfLiteRegistrationExternalSetInit(r, MyCustomOpInit); TfLiteRegistrationExternalSetFree(r, MyCustomOpFree); TfLiteRegistrationExternalSetPrepare(r, MyCustomOpPrepare); TfLiteRegistrationExternalSetInvoke(r, MyCustomOpEval); return r; } const TfLiteRegistrationExternal* MyCustomOpRegistrationExternal() { // Singleton instance, intentionally never destroyed. static const TfLiteRegistrationExternal* my_custom_op = MyCustomOpCreate(); return my_custom_op; } const TfLiteRegistration MyCustomOpRegistration() { static const TfLiteRegistration my_custom_op { .registration_external = MyCustomOpRegistrationExternal(); }; return my_custom_op; }
請注意,註冊不是自動的,應明確呼叫您的 MyCustomOpRegistration
函數(請參閱以下詳細資訊)。雖然標準 BuiltinOpResolver
(可從 :builtin_ops
目標取得)負責內建運算元的註冊,但自訂運算元必須收集在個別的自訂程式庫中。
在 TensorFlow Lite 執行階段中定義核心
我們在 TensorFlow Lite 中使用運算元所需做的就是定義兩個函數 (Prepare
和 Eval
),以及第三個函數來建構 TfLiteRegistrationExternal
C++
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(TfLiteOpaqueTensorData(input)); float* output_data = static_cast (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 TfLiteRegistrationExternal* AtanOpRegistrationExternal() { // Singleton instance, intentionally never destroyed. static const TfLiteRegistrationExternal* atan_op = ()[] { auto* r = TfLiteRegistrationExternalCreate( kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1); TfLiteRegistrationExternalSetPrepare(r, Prepare); TfLiteRegistrationExternalSetInvoke(r, Eval); return r; }; return atan_op; } const TfLiteRegistration AtanOpRegistration() { static const TfLiteRegistration atan_op { .registration_external = AtanOpRegistrationExternal(); }; return atan_op; } } // namespace atan_op
C
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(TfLiteOpaqueTensorData(input)); float* output_data = static_cast (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 TfLiteRegistrationExternal* AtanOpCreate() { TfLiteRegistrationExternal* r = TfLiteRegistrationExternalCreate( kTfLiteBuiltinCustom, "ATAN", /*version=*/ 1); TfLiteRegistrationExternalSetPrepare(r, Prepare); TfLiteRegistrationExternalSetInvoke(r, Eval); return r; } const TfLiteRegistrationExternal* AtanOpRegistrationExternal() { // Singleton instance, intentionally never destroyed. static const TfLiteRegistrationExternal* atan_op = AtanOpCreate(); return atan_op; } const TfLiteRegistration AtanOpRegistration() { static const TfLiteRegistration atan_op { .registration_external = AtanOpRegistrationExternal(); }; return atan_op; }
在初始化 OpResolver
時,將自訂運算元新增至解譯器(請參閱以下範例)。這會向 Tensorflow Lite 註冊運算子,以便 TensorFlow Lite 可以使用新的實作。請注意,TfLiteRegistration
中的最後兩個引數對應於您為自訂運算元定義的 AtanPrepare
和 AtanEval
函數。如果您使用 AtanInit
和 AtanFree
函數來初始化運算元中使用的變數並分別釋放空間,則它們將新增至 TfLiteRegistration
的前兩個引數;在本範例中,這些引數設定為 nullptr
。
向核心程式庫註冊運算子
現在我們需要向核心程式庫註冊運算子。這是透過 OpResolver
完成的。在幕後,解譯器將載入核心程式庫,該程式庫將被指派執行模型中的每個運算子。雖然預設程式庫僅包含內建核心,但可以將其替換/擴充為自訂程式庫運算子。
OpResolver
類別會將運算子程式碼和名稱轉換為實際程式碼,其定義如下:
class OpResolver {
public:
virtual TfLiteRegistration* FindOp(tflite::BuiltinOperator op) const = 0;
virtual TfLiteRegistration* FindOp(const char* op) const = 0;
...
};
請注意,為了向後相容性,此類別使用較舊的具體類型 TfLiteRegistration
而不是不透明類型 TfLiteRegistrationExternal
,但 TfLiteRegistration
結構包含類型為 TfLiteRegistrationExternal*
的 registration_external
欄位。
MutableOpResolver
和 BuiltinOpResolver
類別衍生自 OpResolver
class MutableOpResolver : public OpResolver {
public:
MutableOpResolver(); // Constructs an initially empty op resolver.
void AddBuiltin(tflite::BuiltinOperator op, const TfLiteRegistration* registration) = 0;
void AddCustom(const char* op, const TfLiteRegistration* registration) = 0;
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
,並呼叫 AddCustom
(在您將解譯器傳遞給 InterpreterBuilder
之前)
tflite::ops::builtin::MutableOpResolver resolver;
resolver.AddAll(tflite::ops::builtin::BuiltinOpResolver());
resolver.AddCustom("Atan", AtanOpRegistration());
如果認為內建運算元集太大,則可以根據給定的運算元子集程式碼產生新的 OpResolver
,可能僅限於給定模型中包含的運算元。這相當於 TensorFlow 的選擇性註冊(並且其簡單版本可在 tools 目錄中取得)。
如果您想在 Java 中定義自訂運算子,目前需要在此 jni 程式碼中建構您自己的自訂 JNI 層並編譯您自己的 AAR。同樣地,如果您希望在 Python 中定義這些可用的運算子,您可以將您的註冊放在 Python 包裝函式程式碼中。
請注意,對於支援一組運算而不是單一運算元,可以遵循與上述類似的程序。只需根據需要新增盡可能多的 AddCustom
運算子即可。此外,MutableOpResolver
也允許您使用 AddBuiltin
覆寫內建運算元的實作。
測試及分析您的運算子
若要使用 TensorFlow Lite 基準測試工具分析您的運算元,您可以使用 TensorFlow Lite 的基準模型工具。為了測試目的,您可以透過將適當的 AddCustom
呼叫(如上所示)新增至 register.cc,讓您的 TensorFlow Lite 本機組建知道您的自訂運算元
最佳做法
謹慎最佳化記憶體配置和解除配置。在
Prepare
中配置記憶體比在Invoke
中更有效率,並且在迴圈之前配置記憶體比在每次迭代中都更好。使用暫時張量資料,而不是自行 malloc(請參閱項目 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; }
如果不會浪費太多記憶體,最好使用靜態固定大小陣列(或 Resize 中的預先配置
std::vector
),而不是在每次執行迭代時都使用動態配置std::vector
。避免實例化尚不存在的標準程式庫容器範本,因為它們會影響二進位大小。例如,如果您的運算中需要其他核心中不存在的
std::map
,則使用具有直接索引對應的std::vector
可以運作,同時保持二進位大小較小。查看其他核心使用什麼來獲得見解(或詢問)。檢查 malloc 傳回的記憶體指標。如果此指標為
nullptr
,則不應使用該指標執行任何運算。如果您在函數中 malloc 並有錯誤結束,請在結束前解除配置記憶體。使用
TF_LITE_OPAQUE_ENSURE(context, condition)
檢查特定條件。當使用TF_LITE_OPAQUE_ENSURE
時,您的程式碼不得讓記憶體閒置,也就是說,這些巨集應在配置任何會洩漏的資源之前使用。