Объединение операций TensorFlow

Обзор

На этой странице описаны конструкция и шаги, необходимые для преобразования составных операций в TensorFlow в объединенные операции в LiteRT. Эта инфраструктура является универсальной и поддерживает преобразование любой составной операции в TensorFlow в соответствующую объединенную операцию в LiteRT.

Примером использования этой инфраструктуры является слияние операций TensorFlow RNN с LiteRT, как подробно описано здесь .

Что такое объединенные операции

рисунок

Операции TensorFlow могут быть либо примитивными операциями, например tf.add , либо состоять из других примитивных операций, например tf.einsum . Примитивная операция отображается как один узел в графе TensorFlow, а составная операция представляет собой набор узлов в графе TensorFlow. Выполнение составной операции эквивалентно выполнению каждой из составляющих ее примитивных операций.

Объединенная операция соответствует одной операции, которая включает в себя все вычисления, выполняемые каждой примитивной операцией, в рамках соответствующей составной операции.

Преимущества объединенных операций

Объединенные операции существуют для максимизации производительности лежащих в их основе реализаций ядра за счет оптимизации общих вычислений и уменьшения объема памяти. Это очень ценно, особенно для рабочих нагрузок вывода с малой задержкой и мобильных платформ с ограниченными ресурсами.

Объединенные операции также предоставляют интерфейс более высокого уровня для определения сложных преобразований, таких как квантование, которые в противном случае было бы невозможно или очень сложно выполнить на более детальном уровне.

В LiteRT имеется множество примеров объединенных операций по причинам, изложенным выше. Эти объединенные операции обычно соответствуют составным операциям в исходной программе TensorFlow. Примеры составных операций в TensorFlow, которые реализованы как одна объединенная операция в LiteRT, включают различные операции RNN, такие как однонаправленная и двунаправленная последовательность LSTM, свертка (convolution (conv2d, смещение добавления, relu), полное соединение (matmul, смещение добавления, relu) и другие. В LiteRT квантование LSTM в настоящее время реализовано только в объединенных операциях LSTM.

Проблемы с объединенными операциями

Преобразование составных операций из TensorFlow в объединенные операции в LiteRT — сложная проблема. Это потому, что:

  1. Составные операции представлены в графе TensorFlow как набор примитивных операций без четко определенных границ. Может быть очень сложно идентифицировать (например, посредством сопоставления с образцом) подграф, соответствующий такой составной операции.

  2. Может существовать более одной реализации TensorFlow, предназначенной для объединенной операции LiteRT. Например, в TensorFlow существует множество реализаций LSTM (Keras, Babelfish/lingvo и т. д.), и каждая из них состоит из разных примитивных операций, но все они могут быть преобразованы в одну и ту же объединенную операцию LSTM в LiteRT.

Таким образом, преобразование объединенных операций оказалось весьма сложной задачей.

Оберните составную операцию в tf.function

Во многих случаях некоторая часть модели может быть сопоставлена ​​с одной операцией в TFLite. Это может помочь повысить производительность при написании оптимизированной реализации для конкретных операций. Чтобы иметь возможность создать объединенную операцию в TFLite, определите часть графика, которая представляет объединенную операцию, и оберните ее в tf.function с атрибутом «experimental_implements» в tf.function , который имеет значение атрибута tfl_fusable_op со значением true . Если пользовательская операция принимает атрибуты, передайте их как часть того же «experimental_implements».

Пример,

def get_implements_signature():
  implements_signature = [
    # 'name' will be used as a name for the operation.
    'name: "my_custom_fused_op"',
    # attr "tfl_fusable_op" is required to be set with true value.
    'attr {key: "tfl_fusable_op" value { b: true } }',
    # Example attribute "example_option" that the op accepts.
    'attr {key: "example_option" value { i: %d } }' % 10
  ]
  return ' '.join(implements_signature)

@tf.function(experimental_implements=get_implements_signature())
def my_custom_fused_op(input_1, input_2):
  # An empty function that represents pre/post processing example that
  # is not represented as part of the Tensorflow graph.
  output_1 = tf.constant(0.0, dtype=tf.float32, name='first_output')
  output_2 = tf.constant(0.0, dtype=tf.float32, name='second_output')
  return output_1, output_2

class TestModel(tf.Module):
  def __init__(self):
    super(TestModel, self).__init__()
    self.conv_1 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))
    self.conv_2 = tf.keras.layers.Conv2D(filters=1, kernel_size=(3, 3))

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
      tf.TensorSpec(shape=[1, 28, 28, 3], dtype=tf.float32),
  ])
  def simple_eval(self, input_a, input_b):
    return my_custom_fused_op(self.conv_1(input_a), self.conv_2(input_b))

Обратите внимание, что вам не нужно устанавливать allow_custom_ops для конвертера, поскольку атрибут tfl_fusable_op уже подразумевает это.

Реализуйте пользовательскую операцию и зарегистрируйтесь в TFLite Interpreter.

Реализуйте свою объединенную операцию как пользовательскую операцию TFLite — см. инструкции .

Обратите внимание, что имя для регистрации операции должно быть похоже на имя, указанное в атрибуте name в сигнатуре орудия.

Примером операции в примере является

  TfLiteRegistration reg = {};
  // This name must match the name specified in the implements signature.
  static constexpr char kOpName[] = "my_custom_fused_op";
  reg.custom_name = kOpName;
  reg.prepare = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.invoke = [](TfLiteContext* context, TfLiteNode* node) -> TfLiteStatus {
    // Add your code.
    return kTfLiteOk;
  };
  reg.builtin_code = kTfLiteCustom;
  resolver->AddCustom(kOpName, &reg);

Преобразование из композитного режима в объединенный (расширенный)

Общая архитектура преобразования составных операций TensorFlow в объединенные операции LiteRT приведена ниже:

рисунок

Оберните составную операцию в tf.function

В исходном коде модели TensorFlow определите и абстрагируйте составную операцию в tf.function с аннотацией функции Experiment_implements . См. пример внедрения поиска . Функция определяет интерфейс, и ее аргументы должны использоваться для реализации логики преобразования.

Напишите код преобразования

Код преобразования пишется для интерфейса функции с аннотацией implements . См. пример объединения для внедрения поиска . Концептуально код преобразования заменяет составную реализацию этого интерфейса объединенной.

На этапе подготовки-композитных функций добавьте код преобразования .

В более сложных случаях можно реализовать сложные преобразования операндов составной операции для получения операндов объединенной операции. См. Керас LSTM . код преобразования в качестве примера.

Конвертировать в LiteRT

Используйте API TFLiteConverter.from_saved_model для преобразования в LiteRT.

Под капотом

Теперь мы опишем высокоуровневые детали общей конструкции преобразования в объединенные операции в LiteRT.

Составление операций в TensorFlow

Использование tf.function с атрибутом функции Experiment_implements позволяет пользователям явно составлять новые операции с использованием примитивных операций TensorFlow и указывать интерфейс, который реализует результирующая составная операция. Это очень полезно, поскольку обеспечивает:

  1. Четко определенная граница составной операции в базовом графе TensorFlow.
  2. Явно укажите интерфейс, который реализует эта операция. Аргументы tf.function соответствуют аргументам этого интерфейса.

В качестве примера давайте рассмотрим составную операцию, определенную для реализации встроенного поиска. Это соответствует объединенной операции в LiteRT.

  @tf.function(
        experimental_implements="embedding_lookup")
    def EmbFprop(embs, ids_vec):
      """Embedding forward prop.

      Effectively, it computes:
        num = size of ids_vec
        rets = zeros([num, embedding dim])
        for i in range(num):
          rets[i, :] = embs[ids_vec[i], :]
        return rets

      Args:
        embs: The embedding matrix.
        ids_vec: A vector of int32 embedding ids.

      Returns:
        The result of embedding lookups. A matrix of shape
        [num ids in ids_vec, embedding dims].
      """
      num = tf.shape(ids_vec)[0]
      rets = inplace_ops.empty([num] + emb_shape_suf, py_utils.FPropDtype(p))

      def EmbFpropLoop(i, embs, ids_vec, rets):
        # row_id = ids_vec[i]
        row_id = tf.gather(ids_vec, i)
        # row = embs[row_id]
        row = tf.reshape(tf.gather(embs, row_id), [1] + emb_shape_suf)
        # rets[i] = row
        rets = inplace_ops.alias_inplace_update(rets, [i], row)
        return embs, ids_vec, rets

      _, _, rets = functional_ops.For(
          start=0,
          limit=num,
          delta=1,
          inputs=[embs, ids_vec, rets],
          body=EmbFpropLoop,
          rewrite_with_while=compiled)
      if len(weight_shape) > 2:
        rets = tf.reshape(rets, [num, symbolic.ToStatic(p.embedding_dim)])
      return rets

Заставляя модели использовать составные операции через tf.function , как показано выше, становится возможным построить общую инфраструктуру для идентификации и преобразования таких операций в объединенные операции LiteRT.

Расширение конвертера LiteRT

Конвертер LiteRT, выпущенный ранее в этом году, поддерживал импорт моделей TensorFlow только в виде графика, в котором все переменные заменялись соответствующими постоянными значениями. Это не работает для операций слияния, поскольку в такие графы встроены все функции, поэтому переменные можно превратить в константы.

Чтобы использовать tf.function с функцией experimental_implements во время процесса преобразования, функции необходимо сохранить до более позднего процесса преобразования.

Таким образом, мы реализовали новый рабочий процесс импорта и преобразования моделей TensorFlow в конвертере для поддержки варианта использования составной операции слияния. В частности, добавлены следующие новые функции:

  1. Импорт сохраненных моделей TensorFlow в MLIR
  2. объединение сложных операций
  3. анализ изменчивости переменных
  4. заморозить все переменные, доступные только для чтения

Это позволяет нам выполнять объединение операций, используя функции, представляющие составные операции, до встраивания функций и замораживания переменных.

Реализация операции слияния

Давайте рассмотрим проход операции Fusion более подробно. Этот проход делает следующее:

  1. Прокрутите все функции модуля MLIR.
  2. Если функция имеет атрибут tf._implements, на основе значения атрибута вызывает соответствующую утилиту объединения операций.
  3. Утилита операции слияния работает с операндами и атрибутами функции (которые служат интерфейсом для преобразования) и заменяет тело функции эквивалентным телом функции, содержащим объединенную операцию.
  4. Во многих случаях замененное тело будет содержать операции, отличные от операции слияния. Они соответствуют некоторым статическим преобразованиям операндов функции для получения операндов объединенной операции. Поскольку все эти вычисления могут быть свернуты константами, они не будут присутствовать в экспортированном плоском буфере, где будет существовать только объединенная операция.

Вот фрагмент кода из прохода, показывающий основной рабочий процесс:

void PrepareCompositeFunctionsPass::ConvertTFImplements(FuncOp func,
                                                        StringAttr attr) {
  if (attr.getValue() == "embedding_lookup") {
    func.eraseBody();
    func.addEntryBlock();
    // Convert the composite embedding_lookup function body to a
    // TFLite fused embedding_lookup op.
    ConvertEmbeddedLookupFunc convert_embedded_lookup(func);
    if (failed(convert_embedded_lookup.VerifySignature())) {
      return signalPassFailure();
    }
    convert_embedded_lookup.RewriteFunc();
  } else if (attr.getValue() == mlir::TFL::kKerasLstm) {
     func.eraseBody();
     func.addEntryBlock();
     OpBuilder builder(func.getBody());
     if (failed(ConvertKerasLSTMLayer(func, &builder))) {
       return signalPassFailure();
     }
  } else if (.....) /* Other fusions can plug in here */
}

Вот фрагмент кода, показывающий сопоставление этой составной операции с объединенной операцией в LiteRT, используя эту функцию в качестве интерфейса преобразования. 

void RewriteFunc() {
    Value lookup = func_.getArgument(1);
    Value value = func_.getArgument(0);
    auto output_type = func_.getType().getResult(0);

    OpBuilder builder(func_.getBody());
    auto op = builder.create<mlir::TFL::EmbeddingLookupOp>(
        func_.getLoc(), output_type, lookup, value);

    builder.create<mlir::ReturnOp>(func_.getLoc(), op.getResult());
  }