TensorFlow Lite 運算元版本

本文件說明 TensorFlow Lite 的運算元版本設定架構。運算元版本設定可讓開發人員將新功能和參數新增至現有的運算元。此外,它還保證以下事項

  • 回溯相容性:新的 TensorFlow Lite 實作應處理舊的模型檔案。
  • 前向相容性:舊的 TensorFlow Lite 實作應處理新版轉換器產生新的模型檔案,前提是未使用任何新功能。
  • 前向不相容性偵測:如果舊的 TensorFlow Lite 實作讀取新的模型,其中包含不支援的新版運算元,則應報告錯誤。

範例:將擴張新增至深度可分離卷積

本文件其餘部分將說明 TFLite 中的運算元版本設定,方法是示範如何將擴張參數新增至深度可分離卷積運算。

瞭解擴張並非理解本文件所必需。請注意

  • 將新增 2 個新的整數參數:dilation_width_factordilation_height_factor
  • 不支援擴張的舊深度可分離卷積核心相當於將擴張係數設定為 1。

變更 FlatBuffer 結構描述

若要將新參數新增至運算元,請變更 lite/schema/schema.fbs 中的選項表。

例如,深度可分離卷積的選項表如下所示

table DepthwiseConv2DOptions {
  padding:Padding;
  stride_w:int;
  stride_h:int;
  depth_multiplier:int;
  fused_activation_function:ActivationFunctionType;
}

新增新參數時

  • 新增註解,指出哪個版本支援哪些參數。
  • 當新的實作取得新新增參數的預設值時,其運作方式應與舊的實作完全相同。

新增新參數後的表格會像這樣

table DepthwiseConv2DOptions {
  // Parameters for DepthwiseConv version 1 or above.
  padding:Padding;
  stride_w:int;
  stride_h:int;
  depth_multiplier:int;
  fused_activation_function:ActivationFunctionType;
  // Parameters for DepthwiseConv version 2 or above.
  dilation_w_factor:int = 1;
  dilation_h_factor:int = 1;
}

應針對新的結構描述重新產生檔案 lite/schema/schema_generated.h

變更 C 結構和核心實作

在 TensorFlow Lite 中,核心實作與 FlatBuffer 定義分離。核心會從 lite/c/builtin_op_data.h 中定義的 C 結構讀取參數。

原始的深度可分離卷積參數如下

typedef struct {
  TfLitePadding padding;
  int stride_width;
  int stride_height;
  int depth_multiplier;
  TfLiteFusedActivation activation;
} TfLiteDepthwiseConvParams;

與 FlatBuffer 結構描述相同,新增註解,指出哪個版本開始支援哪些參數。結果如下所示

typedef struct {
  // Parameters for DepthwiseConv version 1 or above.
  TfLitePadding padding;
  int stride_width;
  int stride_height;
  int depth_multiplier;
  TfLiteFusedActivation activation;
  // Parameters for DepthwiseConv version 2 or above.
  int dilation_width_factor;
  int dilation_height_factor;
} TfLiteDepthwiseConvParams;

也請變更核心實作,從 C 結構讀取新新增的參數。此處省略詳細資訊。

變更 FlatBuffer 讀取程式碼

讀取 FlatBuffer 並產生 C 結構的邏輯位於 lite/core/api/flatbuffer_conversions.cc 中。

更新檔案以處理新的參數,如下所示

TfLiteStatus ParseDepthwiseConv2D(const Operator* op,
                                  ErrorReporter* error_reporter,
                                  BuiltinDataAllocator* allocator,
                                  void** builtin_data) {
  CheckParsePointerParams(op, error_reporter, allocator, builtin_data);

  SafeBuiltinDataAllocator safe_allocator(allocator);

  std::unique_ptr<TfLiteDepthwiseConvParams,
                  SafeBuiltinDataAllocator::BuiltinDataDeleter>
      params = safe_allocator.Allocate<TfLiteDepthwiseConvParams>();
  TF_LITE_ENSURE(error_reporter, params != nullptr);

  const DepthwiseConv2DOptions* schema_params =
      op->builtin_options_as_DepthwiseConv2DOptions();

  if (schema_params != nullptr) {
    params->padding = ConvertPadding(schema_params->padding());
    params->stride_width = schema_params->stride_w();
    params->stride_height = schema_params->stride_h();
    params->depth_multiplier = schema_params->depth_multiplier();
    params->activation =
        ConvertActivation(schema_params->fused_activation_function());

    params->dilation_width_factor = schema_params->dilation_w_factor();
    params->dilation_height_factor = schema_params->dilation_h_factor();
  }

  *builtin_data = params.release();
  return kTfLiteOk;
}

此處不需要檢查運算元版本。當新的實作讀取舊的模型檔案 (其中缺少擴張係數) 時,它會使用 1 作為預設值,而新的核心將與舊的核心一致地運作。

變更核心註冊

MutableOpResolver (在 lite/mutable_op_resolver.h 中定義) 提供一些函數來註冊運算元核心。預設的最小和最大版本為 1

void AddBuiltin(tflite::BuiltinOperator op, TfLiteRegistration* registration,
                int min_version = 1, int max_version = 1);
void AddCustom(const char* name, TfLiteRegistration* registration,
               int min_version = 1, int max_version = 1);

內建運算元在 lite/kernels/register.cc 中註冊。在本範例中,我們實作了一個新的運算元核心,可以處理 DepthwiseConv2D 版本 1 和 2,因此我們需要變更此行

AddBuiltin(BuiltinOperator_DEPTHWISE_CONV_2D, Register_DEPTHWISE_CONV_2D());

改為

AddBuiltin(BuiltinOperator_DEPTHWISE_CONV_2D, Register_DEPTHWISE_CONV_2D(),
             /* min_version = */ 1,
             /* max_version = */ 2);

變更 TFLite 運算元版本

下一步是讓 TFLite 填入執行運算元所需的最低版本。在本範例中,這表示

  • 當所有擴張係數均為 1 時,填入版本 = 1。
  • 否則,填入版本 = 2。

修改 lite/tools/versioning/op_version.cc 中運算元的 GetBuiltinOperatorVersion 函數,方法是將新版本新增至 DepthwiseConv2D 的案例

case BuiltinOperator_DEPTHWISE_CONV_2D:
  auto depthwise_conv_params =
      reinterpret_cast<TfLiteDepthwiseConvParams*>(op_sig.builtin_data);
  TFLITE_DCHECK(depthwise_conv_params != nullptr);
  if (depthwise_conv_params->dilation_width_factor != 1 ||
       depthwise_conv_params->dilation_height_factor != 1) {
    return 2;
  }
  return 1;

更新運算元版本對應

最後一個步驟是將新的版本資訊新增至運算元版本對應。此步驟是必要的,因為我們需要根據此版本對應產生模型所需的最低執行階段版本。

若要執行此操作,您需要在 lite/tools/versioning/runtime_version.cc 中新增新的對應項目。

在本範例中,您需要將以下項目新增至 op_version_map

{ {BuiltinOperator_DEPTHWISE_CONV_2D, 2}, %CURRENT_RUNTIME_VERSION%}

其中 %CURRENT_RUNTIME_VERSION% 對應於 tensorflow/core/public/version.h 中定義的目前執行階段版本。

委派實作

TensorFlow Lite 提供委派 API,可將運算元委派給硬體後端。在委派的 Prepare 函數中,檢查每個節點是否支援 Delegation 程式碼的版本。

const int kMaxVersion = 1;
TfLiteNode* node;
TfLiteRegistration* registration = nullptr;
TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration(context, node_index, &node, &registration));

if (registration->version > kMaxVersion) {
  // Reject the node if the version isn't supported.
}

即使委派僅支援版本 1 運算元,也需要這樣做,以便委派在取得較高版本的運算元時可以偵測到不相容性。