Fuzja operacji TensorFlow

Przegląd

Na tej stronie opisujemy projekt i kroki niezbędne do konwertowania operacji złożonych w TensorFlow na operacje skonsolidowane w TensorFlow Lite. Ta infrastruktura jest przeznaczona do użycia ogólnego i obsługuje konwersję dowolnych operacji złożonych w TensorFlow na odpowiadające im operacje unieruchomione w TensorFlow Lite.

Przykładem użycia tej infrastruktury jest fuzja operacji TensorFlow RNN z TensorFlow Lite, jak wyszczególniono tutaj.

Czym są operacje scalone

rysunek

Operacje TensorFlow mogą być operacjami podstawowymi, np. tf.add, albo skomponowane na podstawie innych operacji podstawowych, np. tf.einsum. Operacja podstawowa jest wyświetlana na wykresie TensorFlow jako pojedynczy węzeł, a operacja złożona to zbiór węzłów na wykresie TensorFlow. Wykonanie operacji złożonej jest równoważne z wykonaniem wszystkich cząstkowych operacji na niej składowych.

Operacja uśredniona odpowiada pojedynczej operacji, która obejmuje wszystkie obliczenia wykonywane w ramach każdej operacji podstawowej w ramach odpowiedniej operacji złożonej.

Zalety operacji skonsolidowanych

Operacje scalone mają na celu maksymalizację wydajności bazowych implementacji jądra przez optymalizację ogólnych obliczeń i zmniejszenie ilości pamięci. Jest to bardzo przydatne, zwłaszcza w przypadku zadań wnioskowania o małych opóźnieniach i platform mobilnych z ograniczonymi zasobami.

Operacje łączone mają też interfejs wyższego poziomu do definiowania złożonych przekształceń, takich jak kwantyfikacja, które w innym przypadku byłyby niewykonalne lub bardzo trudne do wykonania na bardziej szczegółowym poziomie.

TensorFlow Lite ma wiele instancji operacji scalonych z podanych wyżej powodów. Te operacje uśrednione odpowiadają zwykle operacjom złożonych w źródłowym programie TensorFlow. Przykłady operacji złożonych w TensorFlow, które są wdrożone jako pojedyncza operacja uśredniona w TensorFlow Lite, obejmują różne operacje RNN, takie jak sekwencja jednokierunkowa i dwukierunkowa LSTM, splot (konwertowanie, odchylenie add, relu), w pełni połączone (matmul, bias add, relu) i nie tylko. W TensorFlow Lite kwantyfikacja LSTM jest obecnie wdrażana tylko w operacjach uśrednionych LSTM.

Wyzwania związane ze zintegrowanymi operacjami

Konwertowanie w TensorFlow Lite operacji złożonych z TensorFlow na operacje łączone nie jest łatwym problemem. Dzieje się tak, ponieważ:

  1. Operacje złożone są przedstawione na wykresie TensorFlow jako zbiór operacji podstawowych bez dobrze zdefiniowanej granicy. Zidentyfikowanie (np. przez dopasowanie do wzorca) wykresu podrzędnego odpowiadającego takiej operacji złożonej może być bardzo trudne.

  2. Może być więcej niż 1 implementacja TensorFlow kierowana na uśrednioną operację TensorFlow Lite. Na przykład w TensorFlow znajduje się wiele implementacji LSTM (Keras, Babelfish/lingvo itp.), a każda z nich składa się z różnych operacji podstawowych, ale wszystkie można jeszcze przekonwertować na tę samą operację usztywnianej LSTM w TensorFlow Lite.

W związku z tym przekształcenie operacji skonsolidowanych stanowi spore wyzwanie.

Opakuj operację złożoną w funkcji tf.function

W wielu przypadkach część części modelu można zmapować na jedną operację wTFLite. Pomaga to zwiększyć wydajność przy pisaniu zoptymalizowanej implementacji pod kątem konkretnych operacji. Aby utworzyć operację uśrednioną w TFLite, znajdź część wykresu, która przedstawia operację skonsolidowaną, i umieść ją w elemencie tf.function z atrybutem „experimental_implements” do elementu tf.function, który ma wartość tfl_fusable_op o wartości true. Jeśli operacja niestandardowa pobiera atrybuty, przekazuje je w ramach tych samych wartości „experimental_implements”.

Przykład:

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))

Pamiętaj, że nie musisz konfigurować atrybutu allow_custom_ops w konwerterze, ponieważ atrybut tfl_fusable_op już to sugeruje.

Wdrażanie działań niestandardowych i rejestrowanie się w usłudze TFLite Interpreter

Zaimplementuj operację uśrednioną jako operację niestandardową TFLite – zapoznaj się z instructions.

Nazwa, za pomocą której rejestrujesz operację, powinna być podobna do nazwy określonej w atrybucie name w podpisie implements.

Przykładem operacji w przykładzie jest

  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);

Konwertuję z operacji złożonej na utrwaloną (zaawansowane)

Ogólna architektura konwertowania operacji złożonych TensorFlow na operacje scalone TensorFlow Lite wygląda poniżej:

rysunek

Opakuj operację złożoną w funkcji tf.function

W kodzie źródłowym modelu TensorFlow zidentyfikuj i wyodrębnij operację złożoną do elementu tf.function z adnotacją experimental_implements. Zobacz przykład wyszukiwania umieszczania. Ta funkcja definiuje interfejs, a jego argumenty powinny być używane do implementacji logiki konwersji.

Zapisywanie kodu konwersji

Kod konwersji jest zapisywany w interfejsie funkcji z adnotacją implements. Zapoznaj się z przykładowym połączeniem z wyszukiwaniem umieszczania. Ogólnie kod konwersji zastępuje implementację złożoną tego interfejsu wersją scaloną.

W przekazywaniu funkcji create-composite-functions umieść wtyczkę do kodu konwersji.

W bardziej zaawansowanych zastosowaniach można wdrożyć złożone przekształcenia operandów operacji złożonej, aby uzyskiwać operandy operacji scalonej. Przykład znajdziesz w sekcji kodu konwersji Keras LSTM.

Konwertuj na TensorFlow Lite

Użyj interfejsu API TFLiteConverter.from_saved_model, aby przekonwertować dane na TensorFlow Lite.

Dla zaawansowanych

Opisaliśmy teraz ogólne szczegóły ogólnego projektu przekształcania w TensorFlow Lite operacji łączonych na operacje łączone.

Operacje tworzenia w TensorFlow

Użycie metody tf.function z atrybutem experimental_implements umożliwia użytkownikom jawne tworzenie nowych operacji za pomocą operacji podstawowych TensorFlow i określanie interfejsu zaimplementowanego przez wynikową operację składaną. Jest to bardzo przydatne, ponieważ zawiera:

  1. Dobrze zdefiniowana granica operacji złożonej na bazowym wykresie TensorFlow.
  2. Wyraźnie określ interfejs zaimplementowany przez tę operację. Argumenty funkcji tf.function odpowiadają argumentom tego interfejsu.

Weźmy jako przykład operację złożona zdefiniowaną w celu zaimplementowania wyszukiwania umieszczania. Jest to mapowane na operację uśrednioną w TensorFlow Lite.

  @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

Dzięki temu, że modele wykorzystują operacje złożone za pomocą funkcji tf.function, tak jak pokazano powyżej, można zbudować ogólną infrastrukturę do identyfikowania i konwertowania takich operacji na operacje skonsolidowane TensorFlow Lite.

Rozszerzanie konwertera TensorFlow Lite

Wprowadzony na początku tego roku konwerter TensorFlow Lite obsługiwał tylko importowanie modeli TensorFlow w formie wykresu ze wszystkimi zmiennymi zastąpionymi odpowiednimi wartościami stałymi. Nie działa to w przypadku łączenia operacji, ponieważ na takich wykresach wszystkie funkcje są wbudowane, tak aby zmienne można było przekształcić w stałe.

Aby korzystać z funkcji tf.function z funkcją experimental_implements w trakcie procesu konwersji, funkcje muszą zostać zachowane do końca tego procesu.

W związku z tym wdrożyliśmy nowy przepływ pracy dotyczący importowania i konwertowania modeli TensorFlow w konwerterze, aby obsługiwać przypadek użycia fuzji operacji złożonych. W szczególności dodaliśmy nowe funkcje:

  1. Importowanie zapisanych modeli TensorFlow do MLIR
  2. operacje złożone z bezpieczników
  3. analiza zmienności zmiennych
  4. zablokować wszystkie zmienne tylko do odczytu

Dzięki temu możemy przeprowadzić fuzję operacji przy użyciu funkcji reprezentujących operacje złożone przed wbudowaniem funkcji i zamrożeniem zmiennych.

Wdrażanie operacji fusion

Przyjrzyjmy się teraz bardziej szczegółowo operacji w ramach fuzji jądrowej. Ta karta ma następujące funkcje:

  1. Zrób pętlę przez wszystkie funkcje w module MLIR.
  2. Jeśli funkcja zawiera atrybut tf._implements zależny od wartości atrybutu, wywołuje odpowiednie narzędzie do operacji Fusion.
  3. Narzędzie operacji fusion działa na operandach i atrybutach funkcji (które pełnią funkcję interfejsu konwersji) i zastępuje treść funkcji jej treścią odpowiadającą jej treści zawierającej operację skonsolidowaną.
  4. W wielu przypadkach zastąpiona treść będzie zawierać operacje inne niż operacja uśredniona. Odpowiadają one niektórym transformacjom statycznym w operandach funkcji w celu uzyskania operandów operacji uśrednionej. Ponieważ obliczenia mogą być stale składane, nie pojawią się w wyeksportowanym płaskim buforze, w którym istniałaby tylko operacja uśredniona.

Oto fragment kodu z karnetu przedstawiający główny przepływ pracy:

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 */
}

Oto fragment kodu pokazujący mapowanie tej operacji złożonej na operację skonsolidowaną w TensorFlow Lite z wykorzystaniem tej funkcji jako interfejsu konwersji.

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());
  }