Fuzja operacji TensorFlow

Omówienie

Na tej stronie opisujemy projekt i kroki wymagane do konwertowania operacji złożonych w TensorFlow do scalonych operacji w LiteRT. Ta infrastruktura jest ogólnego przeznaczenia i obsługuje konwersję dowolnych operacji złożonych w TensorFlow z odpowiednią scaloną operacją w LiteRT.

Przykładem wykorzystania tej infrastruktury jest TensorFlow RNN – operacja fuzji LiteRT znajdziesz tutaj.

Czym są operacje scalone

rysunek

Operacje TensorFlow mogą być operacjami podstawowymi, np. tf.add lub mogą być powstałe na podstawie innych podstawowych operacji, np. tf.einsum. Element podstawowy operacja jest widoczna na wykresie TensorFlow jako pojedynczy węzeł, podczas gdy to zbiór węzłów na grafie TensorFlow. Wykonanie polecenia operacja złożona jest równoznaczna z wykonaniem każdego z jego składowych podstawowych elementów operacji.

Operacja scalona odpowiada pojedynczej operacji, która obejmuje wszystkie wykonywane przez każdą operację podstawową w operacji złożonej.

Zalety operacji scalonych

Operacje scalone mają na celu maksymalizację wydajności bazowego jądra wdrożenia przez optymalizację ogólnych obliczeń i zmniejszenie ilości pamięci śladu Twojej aktywności. Jest to bardzo przydatne, zwłaszcza w przypadku zadań wymagających wnioskowania o małym czasie oczekiwania i platform mobilnych z ograniczeniami zasobów.

Operacje scalone zapewniają również interfejs wyższego poziomu do definiowania złożonych przekształceń takich jak kwantyzacja, które w innym przypadku byłyby niewykonalne lub co jest utrudnione na bardziej szczegółowym poziomie.

LiteRT ma wiele instancji uśrednionych operacji omawianego powyżej. Te scalone operacje zazwyczaj odpowiadają w źródłowym programie TensorFlow. Przykłady operacji złożonych w argumencie TensorFlow, które są zaimplementowane jako pojedyncza operacja scalona w LiteRT obejmują różne operacje RNN, takie jak sekwencja jednokierunkowa i dwukierunkowa LSTM, splot (konw2d, dodanie uprzedzenia, relu), pełne połączenie (matmul, dodanie uprzedzenia, relu) i inne. W LiteRT kwantyzacja LSTM jest obecnie dostępna tylko w ramach połączonych operacji LSTM.

Wyzwania związane z operacjami scalonymi

Konwertuję operacje złożone z TensorFlow na operacje scalone w LiteRT jest trudnym problemem. Dzieje się tak, ponieważ:

  1. Operacje złożone są przedstawione na wykresie TensorFlow jako zbiór w przypadku operacji podstawowych bez dobrze zdefiniowanej granicy. Bardzo często trudne do zidentyfikowania (np. przez dopasowywanie do wzorca) podgrafu odpowiadającej takiej operacji złożonej.

  2. Może być więcej niż jedna implementacja TensorFlow ukierunkowana na uśrednioną Operacja LiteRT. Na przykład istnieje wiele implementacji mniej bezpiecznych aplikacji w TensorFlow (Keras, Babelfish/lingvo itp.), a każdy z nich składa się operacji podstawowych, ale wszystkie można przekonwertować na tę samą operację scalonego LSTM w LiteRT.

W związku z tym przekonwertowanie operacji scalonych okazało się dość trudne.

Umieść operację złożoną w tf.function

W wielu przypadkach część modelu można zmapować na pojedynczą operację TFLite. Może to pomóc w zwiększeniu wydajności podczas pisania zoptymalizowanej implementacji. i konkretnych operacji. Aby móc utworzyć operację scaloną w TFLite, zidentyfikuj część wykresu, która reprezentuje operację ujednoliconą, i opakuj ją tf.function z „eksperymentalne_implementacje” do atrybutu tf.function, który ma atrybut wartość tfl_fusable_op o wartości true. Jeśli operacja niestandardowa a następnie przekazują je w ramach tych samych implementacji „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 ustawiać pola allow_custom_ops w konwerterze jako Atrybut tfl_fusable_op już to sugeruje.

Wdróż niestandardową operację i zarejestruj się za pomocą TFLite Interpreter.

Wdróż operację uśrednioną jako operację TFLite niestandardową – patrz instrukcje.

Nazwa służąca do zarejestrowania operacji powinna być podobna do nazwy określone w atrybucie name w implementacji podpisu.

Przykład dla operacji w tym przykładzie to

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

Konwersja z operacji złożonej na ujednoliconą (zaawansowane)

Ogólna architektura konwertowania operacji złożonych TensorFlow na Operacje scalone LiteRT znajdziesz poniżej:

rysunek

Umieść operację złożoną w tf.function

W kodzie źródłowym modelu TensorFlow znajdź i wyodrębnij element złożony w tf.function za pomocą funkcji experimental_implements adnotacja funkcji. Zobacz przykład wyszukiwania umieszczania. definiuje interfejs, a jego argumenty powinny być używane do implementacji to logika konwersji.

Zapisz kod konwersji

Kod konwersji jest napisany za pomocą interfejsu funkcji ze zdefiniowanym separatorem implements adnotacja. Zobacz przykład fuzji do umieszczania lookup. Kod konwersji zastępuje element złożony z wdrożeniem tego interfejsu.

W ramach karty „przygotowywanie funkcji złożonych” w konwersji w kodzie.

W bardziej zaawansowanych zastosowaniach można wdrożyć złożone przekształcenia operandy operacji złożonej w celu uzyskania operandów scalonego . Zobacz listę Keras LSTM. kodu konwersji.

Przekonwertuj na LiteRT

Użyj TFLiteConverter.from_saved_model Interfejs API do konwersji na LiteRT.

Dla zaawansowanych

W dalszej części artykułu opisujemy ogólne szczegóły projektu konwersji na scalone. i operacji w skali LiteRT.

Tworzenie operacji w TensorFlow

użycie parametru tf.function. z experimental_implements atrybut funkcji pozwala użytkownikom na jawne kompilowanie nowych operacji za pomocą Operacje podstawowe TensorFlow i określ interfejs, którego za pomocą operacji złożonych. Jest to bardzo przydatne, ponieważ zapewnia:

  1. Dobrze zdefiniowana granica operacji złożonej w Wykres TensorFlow.
  2. Wyraźnie określ interfejs implementowany przez tę operację. argumentów funkcji tf.function odpowiadają argumentom tego interfejsu.

Przyjrzyjmy się przykładowej operacji złożonej zdefiniowanej w celu zaimplementowania w przypadku wyszukiwania wektora dystrybucyjnego. Jest to mapowane na operację scaloną w 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

Dzięki możliwości użycia przez modele operacji złożonych za pomocą tf.function jako Jak widać powyżej, możliwe staje się budowanie ogólnej infrastruktury, zidentyfikować i przekonwertować takie operacje na scalone operacje LiteRT;

Rozszerzanie konwertera LiteRT

Konwerter LiteRT, który został wprowadzony na początku tego roku, jest obsługiwany tylko importowanie modeli TensorFlow w postaci grafu, zamieniając wszystkie zmienne na ich odpowiednich wartości stałych. To nie działa w przypadku operacji scalenia, ponieważ Takie wykresy mają wbudowane wszystkie funkcje, dzięki czemu zmienne można przekształcać w stałe.

Aby móc korzystać z tf.function za pomocą funkcji experimental_implements w trakcie procesu konwersji funkcje muszą zostać zachowane na późniejszym etapie procesu konwersji.

W związku z tym wdrożyliśmy nowy przepływ pracy importowania i konwertowania TensorFlow. w konwertocie, aby obsługiwać przypadek użycia operacji złożonej. Nowe, dodane funkcje to:

  1. importowanie zapisanych modeli TensorFlow do MLIR;
  2. połącz operacje złożone
  3. analiza zmienności zmiennej
  4. zablokować wszystkie zmienne tylko do odczytu

Pozwala to przeprowadzić fuzję operacji, używając funkcji reprezentujących przed wbudowywaniem funkcji i blokowaniem zmiennych.

Wdrożenie operacji fuzji

Przyjrzyjmy się bardziej szczegółowo operacji Fusion Pass. Ta karta :

  1. Powtórz wszystkie funkcje w module MLIR.
  2. Jeśli funkcja ma atrybut tf._implements, na podstawie atrybutu wartość, wywołuje odpowiednie narzędzie Fusion.
  3. Narzędzie Fusion Operations działa na operandach funkcji i (które pełnią funkcję interfejsu konwersji) i zastępują o treści funkcji z równoważną treścią funkcji zawierającą operacji uśrednionej.
  4. W wielu przypadkach zastąpiona treść będzie zawierać operacje inne niż operacji uśrednionej. Odpowiadają one niektórym statycznym przekształceniom operandy funkcji w celu uzyskania operandów operacji scalonej. Ponieważ te obliczenia można na bieżąco przenosić, nie można ich znajduje się w wyeksportowanym płaskim buforze, gdzie tylko operacja ujednolicona istnieje.

Oto fragment kodu z karty przedstawiającej 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 ukazujący mapowanie tej operacji złożonej na uśrednioną i wykorzystując ją jako interfejs 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());
  }