Fusioni i funksionimit TensorFlow

Përmbledhje

Kjo faqe përshkruan projektimin dhe hapat e nevojshëm për të konvertuar operacionet e përbëra në TensorFlow në operacione të shkrira në LiteRT. Kjo infrastrukturë është për qëllime të përgjithshme dhe mbështet konvertimin e çdo operacioni të përbërë në TensorFlow në një operacion të shkrirë përkatës në LiteRT.

Një shembull përdorimi i kësaj infrastrukture është bashkimi i operacioneve TensorFlow RNN me LiteRT, siç detajohet këtu .

Cilat janë operacionet e bashkuara

vizatim

Operacionet TensorFlow mund të jenë operacione primitive, p.sh. tf.add , ose mund të përbëhen nga operacione të tjera primitive, p.sh. tf.einsum . Një operacion primitiv shfaqet si një nyje e vetme në grafikun TensorFlow, ndërsa një operacion i përbërë është një koleksion nyjesh në grafikun TensorFlow. Ekzekutimi i një operacioni të përbërë është ekuivalent me ekzekutimin e secilit prej operacioneve primitive përbërëse të tij.

Një operacion i shkrirë korrespondon me një operacion të vetëm që përfshin të gjitha llogaritjet e kryera nga secili operacion primitiv brenda operacionit përkatës të përbërë.

Përfitimet e operacioneve të bashkuara

Operacionet e bashkuara ekzistojnë për të maksimizuar performancën e implementimeve të tyre themelore të bërthamës, duke optimizuar llogaritjen e përgjithshme dhe duke zvogëluar gjurmën e kujtesës. Kjo është shumë e vlefshme, veçanërisht për ngarkesat e punës së nxjerrjes së përfundimeve me vonesë të ulët dhe platformat mobile me burime të kufizuara.

Operacionet e shkrira ofrojnë gjithashtu një ndërfaqe të nivelit më të lartë për të përcaktuar transformime komplekse si kuantizimi, të cilat përndryshe do të ishin të pamundura ose shumë të vështira për t'u bërë në një nivel më të detajuar.

LiteRT ka shumë raste të operacioneve të shkrira për arsyet e artikuluara më sipër. Këto operacione të shkrira zakonisht korrespondojnë me operacionet e përbëra në programin burimor TensorFlow. Shembuj të operacioneve të përbëra në TensorFlow që zbatohen si një operacion i vetëm i shkrirë në LiteRT përfshijnë operacione të ndryshme RNN si LSTM me sekuencë unidireksionale dhe bidireksionale, konvolucion (conv2d, bias add, relu), plotësisht i lidhur (matmul, bias add, relu) dhe më shumë. Në LiteRT, kuantizimi LSTM aktualisht zbatohet vetëm në operacionet e shkrira LSTM.

Sfidat me operacionet e bashkuara

Konvertimi i operacioneve të përbëra nga TensorFlow në operacione të shkrira në LiteRT është një problem i vështirë. Kjo për shkak se:

  1. Operacionet e përbëra përfaqësohen në grafikun TensorFlow si një grup operacionesh primitive pa një kufi të përcaktuar mirë. Mund të jetë shumë e vështirë të identifikohet (p.sh., nëpërmjet përputhjes së modeleve) nëngrafi që korrespondon me një operacion të tillë të përbërë.

  2. Mund të ketë më shumë se një implementim të TensorFlow që synon një operacion të shkrirë LiteRT. Për shembull, ka shumë implementime LSTM në TensorFlow (Keras, Babelfish/lingvo etj.) dhe secila prej tyre përbëhet nga operacione të ndryshme primitive, por të gjitha ato mund të konvertohen në të njëjtin operacion të shkrirë LSTM në LiteRT.

Si i tillë, konvertimi i operacioneve të shkrira ka rezultuar mjaft sfidues.

Mbështillni operacionin e përbërë në një tf.function

Në shumë raste, një pjesë e modelit mund të lidhet me një operacion të vetëm në TFLite. Kjo mund të ndihmojë me performancën kur shkruani një implementim të optimizuar për operacione specifike. Për të qenë në gjendje të krijoni një operacion të shkrirë në TFLite, identifikoni pjesën e grafikut që përfaqëson një operacion të shkrirë dhe mbështilleni atë në një tf.function me atributin "experimental_implements" në një tf.function , i cili ka vlerën e atributit tfl_fusable_op me vlerën true . Nëse operacioni i personalizuar merr atribute, atëherë kalojini ato si pjesë të të njëjtit "experimental_implements".

Shembull,

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

Vini re se nuk keni nevojë të vendosni allow_custom_ops në konvertues pasi atributi tfl_fusable_op e nënkupton këtë tashmë.

Implementoni operacionin e personalizuar dhe regjistrohuni me TFLite Interpreter

Implementoni operacionin tuaj të shkrirë si një operacion të personalizuar TFLite - shihni udhëzimet .

Vini re se emri me të cilin do të regjistrohet operacioni duhet të jetë i ngjashëm me emrin e specifikuar në atributin name në nënshkrimin e implements.

Një shembull për operacionin në shembull është

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

Konvertimi nga operacioni i përbërë në atë të shkrirë (i Avancuar)

Arkitektura e përgjithshme për konvertimin e operacioneve të përbëra TensorFlow në operacione të shkrira LiteRT është më poshtë:

vizatim

Mbështillni operacionin e përbërë në një tf.function

Në kodin burimor të modelit TensorFlow, identifikoni dhe abstraktoni operacionin e përbërë në një tf.function me shënimin e funksionit experimental_implements . Shihni një shembull të kërkimit të ngulitur . Funksioni përcakton ndërfaqen dhe argumentet e tij duhet të përdoren për të zbatuar logjikën e konvertimit.

Shkruaj kodin e konvertimit

Kodi i konvertimit shkruhet sipas ndërfaqes së funksionit me shënimin e implements . Shihni një shembull të fusion për embedding lookup . Konceptualisht, kodi i konvertimit zëvendëson implementimin e përbërë të kësaj ndërfaqeje me atë të fused.

Në pasazhin prepare-composite-functions, shtoni kodin tuaj të konvertimit .

Në përdorime më të avancuara, është e mundur të zbatohen transformime komplekse të operandëve të operacionit të përbërë në mënyrë që të nxirren operandët e operacionit të shkrirë. Shih kodin e konvertimit Keras LSTM si shembull.

Konverto në LiteRT

Përdorni API-n TFLiteConverter.from_saved_model për të konvertuar në LiteRT.

Nën kapuç

Tani përshkruajmë detaje të nivelit të lartë të dizajnit të përgjithshëm në konvertimin në operacione të shkrira në LiteRT.

Operacionet e kompozimit në TensorFlow

Përdorimi i tf.function me atributin e funksionit experimental_implements u lejon përdoruesve të hartojnë në mënyrë eksplicite operacione të reja duke përdorur operacionet primitive të TensorFlow dhe të specifikojnë ndërfaqen që zbaton operacioni i përbërë që rezulton. Kjo është shumë e dobishme pasi ofron:

  1. Një kufi i përcaktuar mirë për operacionin e përbërë në grafikun themelor TensorFlow.
  2. Specifikoni në mënyrë të qartë ndërfaqen që zbaton ky operacion. Argumentet e funksionit tf. korrespondojnë me argumentet e kësaj ndërfaqeje.

Si shembull, le të shqyrtojmë një operacion të përbërë të përcaktuar për të zbatuar kërkimin e ngulitur. Kjo i përshtatet një operacioni të shkrirë në 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

Duke i bërë modelet të përdorin operacione të përbëra nëpërmjet tf.function siç ilustrohet më sipër, bëhet e mundur të ndërtohet një infrastrukturë e përgjithshme për të identifikuar dhe konvertuar operacione të tilla në operacione të shkrira LiteRT.

Zgjerimi i konvertuesit LiteRT

Konvertuesi LiteRT që u publikua më herët këtë vit mbështeti importimin e modeleve TensorFlow vetëm si një grafik me të gjitha variablat e zëvendësuara me vlerat e tyre konstante përkatëse. Kjo nuk funksionon për bashkimin e operacioneve pasi grafikë të tillë i kanë të gjitha funksionet të integruara në mënyrë që variablat të mund të shndërrohen në konstante.

Për të përdorur funksionin tf. me veçorinë experimental_implements gjatë procesit të konvertimit, funksionet duhet të ruhen deri më vonë në procesin e konvertimit.

Si i tillë, ne zbatuam një rrjedhë të re pune për importimin dhe konvertimin e modeleve TensorFlow në konvertues për të mbështetur rastin e përdorimit të bashkimit të operacioneve të përbëra. Në mënyrë specifike, veçoritë e reja të shtuara janë:

  1. Importimi i modeleve të ruajtura të TensorFlow në MLIR
  2. operacionet e përbërë të siguresave
  3. analiza e ndryshueshmërisë së ndryshueshme
  4. ngrini të gjitha variablat vetëm për lexim

Kjo na lejon të kryejmë operacionin e bashkimit duke përdorur funksionet që përfaqësojnë operacionet e përbëra para rreshtimit të funksionit dhe ngrirjes së variablave.

Zbatimi i operacionit të bashkimit

Le ta shqyrtojmë operacionin fusion pass në më shumë detaje. Ky pass bën si në vijim:

  1. Përsërit të gjitha funksionet në modulin MLIR.
  2. Nëse një funksion ka atributin tf._implements, bazuar në vlerën e atributit, thirret programi i duhur i bashkimit të operacioneve.
  3. Shërbimi i operacionit të bashkimit operon mbi operandët dhe atributet e funksionit (të cilat shërbejnë si ndërfaqe për konvertimin) dhe zëvendëson trupin e funksionit me një trup funksioni ekuivalent që përmban operacionin e bashkuar.
  4. Në shumë raste, trupi i zëvendësuar do të përmbajë operacione të tjera përveç operacionit të shkrirë. Këto korrespondojnë me disa transformime statike në operandët e funksionit për të marrë operandët e operacionit të shkrirë. Meqenëse këto llogaritje mund të jenë të gjitha konstante të palosura, ato nuk do të ishin të pranishme në flatbuffer-in e eksportuar ku do të ekzistonte vetëm operacioni i shkrirë.

Ja një fragment kodi nga kalimi që tregon rrjedhën kryesore të punës:

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

Ja një fragment kodi që tregon lidhjen e këtij operacioni të përbërë me një operacion të shkrirë në LiteRT duke përdorur funksionin si një ndërfaqe konvertimi.

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