Версии операторов 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. Ядра считывают параметр из структур C, определенных в lite/c/builtin_op_data.h .

Исходный параметр глубинной свертки выглядит следующим образом:

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 . В этом примере мы реализовали новое ядро ​​op, которое может обрабатывать 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.

Измените функцию GetBuiltinOperatorVersion для оператора в lite/tools/versioning/op_version.cc добавив новую версию в случае 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, чтобы делегирование могло обнаружить несовместимость при получении операции более высокой версии.