Fuzja operacji TensorFlow

Przegląd

Na tej stronie opisano projekt i kroki potrzebne do konwersji operacji złożonych w TensorFlow na operacje połączone w TensorFlow Lite. Infrastruktura ta ma charakter ogólny i obsługuje konwersję dowolnej operacji złożonej w TensorFlow na odpowiednią połączoną operację w TensorFlow Lite.

Przykładowym zastosowaniem tej infrastruktury jest połączenie operacji TensorFlow RNN z TensorFlow Lite, jak opisano szczegółowo tutaj .

Co to są operacje połączone

rysunek

Operacje TensorFlow mogą być prymitywnymi operacjami, np. tf.add , lub mogą składać się z innych prymitywnych operacji, np. tf.einsum . Operacja pierwotna jest wyświetlana jako pojedynczy węzeł na wykresie TensorFlow, podczas gdy operacja złożona jest zbiorem węzłów na wykresie TensorFlow. Wykonanie operacji złożonej jest równoznaczne z wykonaniem każdej z jej składowych operacji pierwotnych.

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

Korzyści z operacji połączonych

Operacje fuzyjne mają na celu maksymalizację wydajności 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.

Połączone operacje zapewniają również interfejs wyższego poziomu do definiowania złożonych transformacji, takich jak kwantyzacja, co w innym przypadku byłoby niewykonalne lub bardzo trudne do wykonania na bardziej szczegółowym poziomie.

TensorFlow Lite ma wiele przypadków połączonych operacji z powodów opisanych powyżej. Te połączone operacje zazwyczaj 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, splot (conv2d, bias add, relu), w pełni połączone (matmul, bias add, relu) i więcej . W TensorFlow Lite kwantyzacja LSTM jest obecnie wdrażana tylko w połączonych operacjach LSTM.

Wyzwania związane z operacjami połączonymi

Konwersja operacji złożonych z TensorFlow na operacje połączone w TensorFlow Lite jest trudnym problemem. To dlatego, że:

  1. Operacje złożone są reprezentowane na wykresie TensorFlow jako zbiór operacji pierwotnych bez dobrze określonej granicy. Zidentyfikowanie (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 prymitywnych operacji, ale wszystkie nadal można przekonwertować na tę samą połączoną operację LSTM w TensorFlow Lite.

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

Zawiń 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 zwiększeniu wydajności podczas pisania zoptymalizowanej implementacji dla określonych operacji. Aby móc utworzyć połączoną operację w TFLite, zidentyfikuj część wykresu reprezentującą połączoną operację i zawiń ją w tf.function z atrybutem „experimental_implements” do tf.function , która 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 tłumaczu TFLite

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

Zauważ, że nazwa, pod którą chcesz zarejestrować operację, powinna być podobna do nazwy określonej w atrybucie name w sygnaturze implements.

Przykładem operacji w tym 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);

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

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

rysunek

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

W kodzie źródłowym modelu TensorFlow zidentyfikuj i wyabstrahuj operację złożoną do tf.function z adnotacją funkcji eksperymentalnej_implements . Zobacz przykład wyszukiwania osadzania . Funkcja definiuje interfejs i jej argumenty służą do implementacji logiki konwersji.

Napisz kod konwersji

Kod konwersji jest zapisywany w interfejsie funkcji z adnotacją implements . Zobacz przykładową fuzję dotyczącą wyszukiwania osadzania . Koncepcyjnie kod konwersji zastępuje złożoną implementację tego interfejsu połączonym.

W przekazie funkcji przygotowania złożonego dodaj wtyczkę do kodu konwersji .

W bardziej zaawansowanych zastosowaniach możliwe jest wdrożenie 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 dotyczące konwersji na operacje połączone w TensorFlow Lite.

Operacje komponowania w TensorFlow

Użycie funkcji tf.function z atrybutem funkcji eksperymentalnej_implements pozwala użytkownikom jawnie tworzyć nowe operacje przy użyciu prymitywnych operacji TensorFlow i określać 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 implementowany przez tę operację. Argumenty tf.function odpowiadają argumentom tego interfejsu.

Jako przykład rozważmy operację złożoną zdefiniowaną w celu implementacji wyszukiwania osadzania. Odwzorowuje to połączoną operację 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 korzystają z operacji złożonych za pośrednictwem tf.function , jak pokazano powyżej, możliwe staje się zbudowanie ogólnej infrastruktury umożliwiającej identyfikację i konwersję takich operacji na połączone operacje TensorFlow Lite.

Rozszerzenie konwertera TensorFlow Lite

Konwerter TensorFlow Lite wydany na początku tego roku umożliwiał jedynie importowanie modeli TensorFlow w postaci wykresu, w którym wszystkie zmienne zastąpiono odpowiadającymi im stałymi wartościami. Nie działa to w przypadku fuzji operacji, ponieważ na takich wykresach wszystkie funkcje są wstawione, dzięki czemu zmienne można zamienić na stałe.

Aby podczas procesu konwersji wykorzystać tf.function z funkcją experimental_implements , funkcje muszą zostać zachowane do późniejszej części procesu 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 złożonych. W szczególności dodane nowe funkcje to:

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

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

Wdrażanie fuzji operacji

Przyjrzyjmy się bardziej szczegółowo operacji fuzji. To przejście wykonuje następujące czynności:

  1. Przejdź przez wszystkie funkcje modułu MLIR.
  2. Jeśli funkcja ma atrybut tf._implements, na podstawie wartości atrybutu wywołuje odpowiednie narzędzie scalające operacje.
  3. Narzędzie Operation Fusion działa na operandach i atrybutach funkcji (które służą jako interfejs do konwersji) i zastępuje treść funkcji równoważną treścią funkcji zawierającą połączoną operację.
  4. W wielu przypadkach zastąpiona bryła będzie zawierać operacje inne niż operacja stopiona. Odpowiadają one pewnym statycznym przekształceniom operandów funkcji w celu uzyskania operandów połączonej operacji. Ponieważ wszystkie te obliczenia można stale składać, nie będą one obecne w wyeksportowanym buforze płaskim, w którym istniałaby tylko operacja łączenia.

Oto fragment kodu z przepustki 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 operacji złożonej na operację połączoną w TensorFlow Lite wykorzystującą 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());
  }