Fuzja operacyjna TensorFlow, Fuzja operacyjna TensorFlow

Zadbaj o dobrą organizację dzięki kolekcji Zapisuj i kategoryzuj treści zgodnie ze swoimi preferencjami.

Przegląd

Ta strona opisuje projekt i kroki potrzebne do konwersji operacji złożonych w TensorFlow na operacje połączone w TensorFlow Lite. Infrastruktura ta jest przeznaczona do celów ogólnych i obsługuje konwersję dowolnej złożonej operacji w TensorFlow do odpowiedniej połączonej operacji w TensorFlow Lite.

Przykładem wykorzystania tej infrastruktury jest fuzja operacji TensorFlow RNN z TensorFlow Lite, jak opisano tutaj .

Co to są operacje połączone?

rysunek

Operacje TensorFlow mogą być albo operacjami prymitywnymi, np. tf.add , albo mogą składać się z innych operacji podstawowych, np. tf.einsum . Operacja pierwotna jest wyświetlana jako pojedynczy węzeł na grafie TensorFlow, podczas gdy operacja złożona jest zbiorem węzłów na grafie TensorFlow. Wykonywanie operacji złożonej jest równoważne wykonaniu każdej z jej składowych operacji pierwotnych.

Operacja połączona odpowiada pojedynczej operacji, która obejmuje wszystkie obliczenia wykonywane przez każdą pierwotną operację w ramach odpowiedniej operacji złożonej.

Korzyści z operacji łączonych

Operacje połączone istnieją w celu maksymalizacji wydajności ich podstawowych implementacji jądra poprzez optymalizację ogólnych obliczeń i zmniejszenie zużycia pamięci. Jest to bardzo cenne, zwłaszcza w przypadku obciążeń wnioskowania o małych opóźnieniach i platform mobilnych o ograniczonych zasobach.

Operacje połączone 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 bardzo trudne do wykonania na bardziej szczegółowym poziomie.

TensorFlow Lite ma wiele przypadków połączonych operacji z powodów wymienionych powyżej. Te połączone operacje zwykle odpowiadają operacjom złożonym w źródłowym programie TensorFlow. Przykłady operacji złożonych w TensorFlow, które są zaimplementowane jako pojedyncza operacja połączona w TensorFlow Lite, obejmują różne operacje RNN, takie jak sekwencja jednokierunkowa i dwukierunkowa LSTM, konwolucja (conv2d, bias add, relu), w pełni połączone (matmul, bias add, relu) i inne . W TensorFlow Lite kwantyzacja LSTM jest obecnie zaimplementowana tylko w połączonych operacjach LSTM.

Wyzwania związane z operacjami połączonymi

Konwersja operacji złożonych z TensorFlow do operacji połączonych w TensorFlow Lite to trudny problem. To dlatego, że:

  1. Operacje złożone są reprezentowane na wykresie TensorFlow jako zbiór operacji pierwotnych bez dobrze zdefiniowanej granicy. Identyfikacja (np. poprzez dopasowanie wzorca) podwykresu odpowiadającego takiej złożonej operacji może być bardzo trudne.

  2. Może istnieć więcej niż jedna implementacja TensorFlow ukierunkowana na połączoną operację TensorFlow Lite. Na przykład istnieje wiele implementacji LSTM w TensorFlow (Keras, Babelfish/lingvo itp.), a każda z nich składa się z różnych operacji pierwotnych, ale wszystkie mogą zostać przekonwertowane do tej samej połączonej operacji LSTM w TensorFlow Lite.

W związku z tym konwersja operacji połączonych okazała się dość trudna.

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

W wielu przypadkach część modelu można zmapować do pojedynczej operacji w TFLite. Może to pomóc w poprawie wydajności podczas pisania zoptymalizowanej implementacji dla określonych operacji. Aby móc utworzyć operację połączoną w TFLite, zidentyfikuj część wykresu, która reprezentuje operację połączoną i zawiń ją w tf.function z atrybutem „experimental_implements” do tf.function , który ma wartość atrybutu tfl_fusable_op z wartością true . Jeśli operacja niestandardowa przyjmuje atrybuty, przekaż je jako część tego samego „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))

Zauważ, że nie musisz ustawiać allow_custom_ops na konwerterze, ponieważ atrybut tfl_fusable_op już to sugeruje.

Zaimplementuj niestandardową operację i zarejestruj się w TFLite Interpreter

Zaimplementuj swoją operację połączoną jako operację niestandardową TFLite — patrz instrukcje .

Zauważ, że nazwa do zarejestrowania operacji powinna być podobna do nazwy określonej w atrybucie name w podpisie implementacji.

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 coder.
    return kTfLiteOk;
  };
  reg.builtin_code = kTfLiteCustom;
  resolver->AddCustom(kOpName, &reg);

Konwersja z operacji kompozytowej do operacji stopionej (zaawansowane)

Ogólna architektura konwersji operacji złożonych TensorFlow na operacje połączone TensorFlow Lite jest poniżej:

rysunek

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

W kodzie źródłowym modelu TensorFlow zidentyfikuj i wydziel operację złożoną do funkcji tf.function z adnotacją funkcji eksperymentalnej_implements . Zobacz przykład osadzania odnośnika . Funkcja definiuje interfejs i jej argumenty powinny być użyte do implementacji logiki konwersji.

Napisz kod konwersji

Kod konwersji jest pisany na interfejsie funkcji z adnotacją o implements . Zobacz przykładowe połączenie wyszukiwania osadzania . Koncepcyjnie kod konwersji zastępuje złożoną implementację tego interfejsu przez fuzję.

W przejściu funkcji przygotowania-złożonych, wstaw swój kod konwersji .

W bardziej zaawansowanych zastosowaniach możliwe jest zaimplementowanie złożonych transformacji operandów operacji złożonej w celu uzyskania operandów operacji połączonej. Zobacz Keras LSTM . kod konwersji jako przykład.

Konwertuj na TensorFlow Lite

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

Pod maską

Opiszemy teraz szczegółowe szczegóły ogólnego projektu konwersji na operacje połączone w TensorFlow Lite.

Operacje komponowania w TensorFlow

Użycie tf.function z atrybutem funkcji exclusive_implements pozwala użytkownikom jawnie komponować nowe operacje przy użyciu operacji prymitywnych TensorFlow i określić interfejs, który implementuje wynikowa operacja złożona. Jest to bardzo przydatne, ponieważ zapewnia:

  1. Dobrze zdefiniowana granica operacji złożonej na podstawowym wykresie TensorFlow.
  2. Jawnie określ interfejs, który implementuje ta operacja. Argumenty funkcji tf.function odpowiadają argumentom tego interfejsu.

Jako przykład rozważmy operację złożoną zdefiniowaną do implementacji wyszukiwania osadzania. Odwzorowuje to działanie połączone 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

Sprawiając, że modele wykorzystują złożone operacje za pośrednictwem tf.function , jak pokazano powyżej, możliwe staje się zbudowanie ogólnej infrastruktury do identyfikowania i przekształcania takich operacji w połączone operacje TensorFlow Lite.

Rozszerzenie konwertera TensorFlow Lite

Konwerter TensorFlow Lite, który został wydany na początku tego roku, obsługiwał tylko importowanie modeli TensorFlow w postaci wykresu ze wszystkimi zmiennymi zastąpionymi odpowiadającymi im wartościami stałymi. Nie działa to w przypadku fuzji operacji, ponieważ takie wykresy mają wszystkie funkcje wbudowane, dzięki czemu zmienne można przekształcić w stałe.

Aby wykorzystać funkcję tf.function z funkcją experimental_implements podczas procesu konwersji, funkcje muszą zostać zachowane na później w procesie konwersji.

W związku z tym wdrożyliśmy nowy przepływ pracy polegający na importowaniu i konwertowaniu modeli TensorFlow w konwerterze w celu obsługi przypadku użycia fuzji operacji kompozytowych. W szczególności dodane nowe funkcje to:

  1. Importowanie zapisanych modeli TensorFlow do MLIR
  2. operacje na kompozytach bezpiecznikowych
  3. zmienna analiza mutacji
  4. zamrozić wszystkie zmienne tylko do odczytu

To pozwala nam wykonać fuzję operacji za pomocą funkcji reprezentujących operacje złożone przed inliningiem funkcji i zmiennym zamrażaniem.

Wdrażanie operacji fuzji

Przyjrzyjmy się bardziej szczegółowo przebiegowi operacji fuzji. Ta przepustka umożliwia:

  1. Przejdź przez wszystkie funkcje w module MLIR.
  2. Jeśli funkcja ma atrybut tf._implements, na podstawie wartości atrybutu wywołuje odpowiednie narzędzie fuzji operacji.
  3. Narzędzie scalania operacji operuje na operandach i atrybutach funkcji (które służą jako interfejs do konwersji) i zastępuje ciało funkcji równoważnym ciałem funkcji zawierającym operację połączoną.
  4. W wielu przypadkach wymieniony korpus będzie zawierał operacje inne niż operacja stopiona. Odpowiadają one pewnym statycznym przekształceniom operandów funkcji w celu uzyskania operandów operacji połączonej. Ponieważ wszystkie te obliczenia mogą być stale składane, nie byłyby obecne w wyeksportowanym płaskim buforze, gdzie istniałaby tylko operacja łączona.

Oto fragment kodu z karty pokazują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 złożonej operacji do połączonej operacji w TensorFlow Lite, wykorzystującej tę funkcję 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());
  }