TensorFlow operasyon füzyonu

Genel Bakış

Bu sayfada TensorFlow'daki kompozit işlemleri TensorFlow Lite'taki birleştirilmiş işlemlere dönüştürmek için gereken tasarım ve adımlar açıklanmaktadır. Bu altyapı genel amaçlıdır ve TensorFlow'daki herhangi bir kompozit işlemin TensorFlow Lite'ta karşılık gelen birleştirilmiş işleme dönüştürülmesini destekler.

Bu altyapının örnek bir kullanımı, burada ayrıntılarıyla anlatıldığı gibi TensorFlow RNN işleminin TensorFlow Lite ile birleştirilmesidir.

Birleştirilmiş işlemler nelerdir

çizim

TensorFlow işlemleri, tf.add gibi ilkel işlemler olabilir veya tf.einsum gibi diğer ilkel işlemlerden oluşabilir. İlkel bir işlem, TensorFlow grafiğinde tek bir düğüm olarak görünürken, bileşik bir işlem, TensorFlow grafiğinde bir düğüm koleksiyonudur. Bileşik bir işlemin yürütülmesi, onu oluşturan temel işlemlerin her birinin yürütülmesine eşdeğerdir.

Birleştirilmiş bir işlem, karşılık gelen bileşik işlem içindeki her bir ilkel işlem tarafından gerçekleştirilen tüm hesaplamaları kapsayan tek bir işleme karşılık gelir.

Birleştirilmiş operasyonların faydaları

Birleştirilmiş işlemler, genel hesaplamayı optimize ederek ve bellek ayak izini azaltarak, temeldeki çekirdek uygulamalarının performansını en üst düzeye çıkarmak için mevcuttur. Bu, özellikle düşük gecikmeli çıkarım iş yükleri ve kaynak kısıtlı mobil platformlar için çok değerlidir.

Birleştirilmiş işlemler aynı zamanda niceleme gibi karmaşık dönüşümleri tanımlamak için daha yüksek düzeyde bir arayüz sağlar; aksi takdirde daha ayrıntılı bir düzeyde yapılması mümkün olmaz veya çok zor olur.

TensorFlow Lite, yukarıda belirtilen nedenlerden dolayı birçok birleşik işlem örneğine sahiptir. Bu birleştirilmiş işlemler tipik olarak kaynak TensorFlow programındaki bileşik işlemlere karşılık gelir. TensorFlow Lite'ta tek bir birleştirilmiş işlem olarak uygulanan TensorFlow'daki bileşik operasyon örnekleri arasında Tek Yönlü ve Çift Yönlü dizi LSTM, evrişim (conv2d,bias add, relu), tamamen bağlı (matmul,bias add, relu) ve daha fazlası gibi çeşitli RNN işlemleri yer alır. . TensorFlow Lite'ta LSTM nicelemesi şu anda yalnızca birleştirilmiş LSTM işlemlerinde uygulanmaktadır.

Birleştirilmiş operasyonlarla ilgili zorluklar

Bileşik işlemleri TensorFlow'dan TensorFlow Lite'ta birleştirilmiş işlemlere dönüştürmek zor bir sorundur. Bunun nedeni ise:

  1. Bileşik işlemler, TensorFlow grafiğinde, iyi tanımlanmış bir sınıra sahip olmayan bir dizi temel işlem olarak temsil edilir. Böyle bir bileşik işleme karşılık gelen alt grafiğin belirlenmesi (örn. desen eşleştirme yoluyla) çok zor olabilir.

  2. Birleştirilmiş bir TensorFlow Lite işlemini hedefleyen birden fazla TensorFlow uygulaması olabilir. Örneğin, TensorFlow'da birçok LSTM uygulaması vardır (Keras, Babelfish/lingvo vb.) ve bunların her biri farklı ilkel işlemlerden oluşur ancak hepsi yine de TensorFlow Lite'ta aynı birleştirilmiş LSTM işlemine dönüştürülebilir.

Bu nedenle, birleştirilmiş operasyonların dönüştürülmesinin oldukça zorlu olduğu kanıtlanmıştır.

Bileşik işlemi bir tf.function sarın

Çoğu durumda modelin bir kısmı TFLite'da tek bir işleme eşlenebilir. Bu, belirli işlemler için optimize edilmiş bir uygulama yazarken performansa yardımcı olabilir. TFLite'da birleştirilmiş bir işlem oluşturabilmek için, grafiğin birleştirilmiş bir işlemi temsil eden kısmını tanımlayın ve bunu "experimental_implements" özniteliğine sahip bir tf.function içine, true değerine sahip tfl_fusable_op öznitelik değerine sahip bir tf.function öğesine sarın. Özel işlem nitelikleri alıyorsa, bunları aynı "experimental_implements" parçası olarak iletin.

Örnek,

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

tfl_fusable_op özelliği zaten bunu ima ettiğinden dönüştürücüde allow_custom_ops ayarlamanıza gerek olmadığını unutmayın.

Özel operasyon uygulayın ve TFLite Interpreter'a kaydolun

Birleştirilmiş işleminizi bir TFLite Özel işlemi olarak uygulayın - talimatlara bakın.

Op'un kaydedileceği ismin, uygulama imzasındaki name özelliğinde belirtilen isme benzer olması gerektiğini unutmayın.

Örnekteki op için bir örnek:

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

Kompozitten kaynaşmış çalışmaya dönüştürme (Gelişmiş)

TensorFlow kompozit işlemlerini TensorFlow Lite birleşik işlemlerine dönüştürmeye yönelik genel mimari aşağıdadır:

çizim

Bileşik işlemi bir tf.function sarın

TensorFlow modelinin kaynak kodunda, birleşik işlemi deneysel_implements işlevi ek açıklamasıyla tanımlayın ve bir tf.function işlevine soyutlayın. Aramayı yerleştirme örneğine bakın. İşlev, arayüzü tanımlar ve onun argümanları, dönüştürme mantığını uygulamak için kullanılmalıdır.

Dönüşüm kodunu yaz

Dönüşüm kodu, fonksiyonun arayüzüne göre implements açıklamasıyla yazılır. Aramayı gömmek için örnek bir füzyona bakın. Kavramsal olarak, dönüşüm kodu bu arayüzün bileşik uygulamasını birleştirilmiş olanla değiştirir.

Hazırlık-bileşik işlevler geçişinde dönüşüm kodunuzu ekleyin.

Daha gelişmiş kullanımlarda, birleştirilmiş işlemin işlenenlerini türetmek amacıyla bileşik işlemin işlenenlerinin karmaşık dönüşümlerini uygulamak mümkündür. Keras LSTM'ye bakın. örnek olarak dönüşüm kodu.

TensorFlow Lite'a dönüştürün

TensorFlow Lite'a dönüştürmek için TFLiteConverter.from_saved_model API'sini kullanın.

Kaputun altında

Artık TensorFlow Lite'ta birleştirilmiş işlemlere dönüştürmede genel tasarımın üst düzey ayrıntılarını açıklıyoruz.

TensorFlow'da işlemler oluşturma

tf.function işlevinin deneysel_implements işlev özelliğiyle kullanılması, kullanıcıların TensorFlow temel işlemlerini kullanarak açıkça yeni işlemler oluşturmasına ve sonuçta ortaya çıkan bileşik işlemin uyguladığı arayüzü belirtmesine olanak tanır. Bu, aşağıdakileri sağladığı için çok faydalıdır:

  1. Temel TensorFlow grafiğindeki bileşik işlem için iyi tanımlanmış bir sınır.
  2. Bu işlemin uyguladığı arayüzü açıkça belirtin. tf.function argümanları bu arayüzün argümanlarına karşılık gelir.

Örnek olarak, gömme aramasını uygulamak için tanımlanmış bir bileşik işlemi ele alalım. Bu, TensorFlow Lite'ta birleştirilmiş bir işlemle eşleşir.

  @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

Modellerin yukarıda gösterildiği gibi tf.function aracılığıyla bileşik operasyonlar kullanmasını sağlayarak, bu tür operasyonları tanımlamak ve birleştirilmiş TensorFlow Lite operasyonlarına dönüştürmek için genel bir altyapı oluşturmak mümkün hale gelir.

TensorFlow Lite dönüştürücünün genişletilmesi

Bu yılın başlarında piyasaya sürülen TensorFlow Lite dönüştürücü, yalnızca TensorFlow modellerinin, tüm değişkenlerin karşılık gelen sabit değerleriyle değiştirildiği bir grafik olarak içe aktarılmasını destekliyordu. Bu tür grafiklerde tüm işlevler satır içi olduğundan, değişkenler sabitlere dönüştürülebildiğinden, bu işlem füzyon için işe yaramaz.

Dönüştürme işlemi sırasında experimental_implements özelliği ile tf.function yararlanmak için, işlevlerin dönüştürme işleminin ilerleyen aşamalarına kadar korunması gerekir.

Bu nedenle, kompozit operasyon füzyon kullanım durumunu desteklemek için dönüştürücüde TensorFlow modellerinin içe aktarılması ve dönüştürülmesine yönelik yeni bir iş akışı uyguladık. Özellikle eklenen yeni özellikler şunlardır:

  1. TensorFlow kayıtlı modellerini MLIR'e aktarma
  2. sigorta kompozit işlemleri
  3. değişken değişkenlik analizi
  4. tüm salt okunur değişkenleri dondur

Bu, fonksiyon satır içi ve değişken dondurma öncesinde bileşik işlemleri temsil eden fonksiyonları kullanarak işlem füzyonu gerçekleştirmemize olanak tanır.

Operasyon füzyonunun uygulanması

Şimdi füzyon geçişi operasyonuna daha detaylı bakalım. Bu geçiş aşağıdakileri yapar:

  1. MLIR modülündeki tüm işlevler arasında geçiş yapın.
  2. Bir işlev, öznitelik değerine bağlı olarak tf._implements özniteliğine sahipse, uygun işlem birleştirme yardımcı programını çağırır.
  3. Operasyon birleştirme yardımcı programı, fonksiyonun işlenenleri ve nitelikleri (dönüştürme için arayüz görevi gören) üzerinde çalışır ve fonksiyonun gövdesini, birleştirilmiş işlemi içeren eşdeğer bir fonksiyon gövdesiyle değiştirir.
  4. Çoğu durumda değiştirilen gövde, birleştirme işlemi dışında başka işlemler de içerecektir. Bunlar, birleştirilmiş işlemin işlenenlerini elde etmek için işlevin işlenenleri üzerindeki bazı statik dönüşümlere karşılık gelir. Bu hesaplamaların tümü sürekli olarak katlanabilir olduğundan, yalnızca birleştirilmiş işlemin mevcut olduğu dışa aktarılan düz arabellekte mevcut olmazlar.

Ana iş akışını gösteren geçişteki kod pasajı:

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

Burada, bu bileşik işlemi TensorFlow Lite'ta bir dönüşüm arayüzü olarak işlevden yararlanan birleştirilmiş bir işlemle eşlemeyi gösteren kod pasajı bulunmaktadır.

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