TensorFlow Lite 运算符版本

本文档介绍了 TensorFlow Lite 的运算符版本控制方案。运算符版本控制使开发人员能够向现有运算符添加新功能和参数。此外，它保证以下内容

• 向后兼容性：新的 TensorFlow Lite 实现应处理旧的模型文件。
• 向前兼容性：旧的 TensorFlow Lite 实现应处理由新版本转换器生成的新模型文件，只要没有使用新功能。
• 向前不兼容性检测：如果旧的 TensorFlow Lite 实现读取包含不支持的新版本运算符的新模型，则应报告错误。

示例：向深度可分离卷积添加膨胀

本文档的其余部分通过展示如何向深度可分离卷积运算添加膨胀参数来解释 TFLite 中的运算符版本控制。

无需了解膨胀即可理解本文档。请注意

• 将添加 2 个新的整数参数：dilation_width_factor 和 dilation_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 函数中，检查委托代码中每个节点的版本是否受支持。

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 操作，这也是必需的，因此委托可以在获取更高版本的操作时检测到不兼容性。