Fusioni i funksionimit TensorFlow

Vështrim i përgjithshëm

Kjo faqe përshkruan dizajnin dhe hapat e nevojshëm për të kthyer operacionet e përbëra në TensorFlow në operacione të bashkuara në LiteRT. Kjo infrastrukturë është me qëllim të përgjithshëm dhe mbështet konvertimin e çdo operacioni të përbërë në TensorFlow në një operacion përkatës të shkrirë në LiteRT.

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

Cilat janë operacionet e bashkuara

vizatim

Operacionet TensorFlow mund të jenë ose ops 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ë i barabartë me ekzekutimin e secilit prej operacioneve primitive përbërëse të tij.

Një operacion i bashkuar korrespondon me një operacion të vetëm që përfshin të gjithë llogaritjen e kryer nga çdo 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 zbatimeve të tyre themelore të kernelit, duke optimizuar llogaritjen e përgjithshme dhe duke reduktuar gjurmën e kujtesës. Kjo është shumë e vlefshme, veçanërisht për ngarkesat e punës me konkluzion me vonesë të ulët dhe platformat celulare me burime të kufizuara.

Operacionet e bashkuara sigurojnë gjithashtu një ndërfaqe të nivelit më të lartë për të përcaktuar transformimet komplekse si kuantizimi, i cili përndryshe do të ishte i pamundur ose shumë i vështirë për t'u bërë në një nivel më të grimcuar.

LiteRT ka shumë raste të operacioneve të bashkuara për arsyet e artikuluara më sipër. Këto operacione të bashkuara 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 sekuenca njëdrejtimëshe dhe dydrejtimëshe LSTM, konvolucioni (conv2d, bias add, relu), i lidhur plotësisht (matmul, bias add, relu) dhe më shumë. Në LiteRT, kuantizimi LSTM aktualisht zbatohet vetëm në operacionet e bashkuara LSTM.

Sfidat me operacionet e bashkuara

Konvertimi i operacioneve të përbëra nga TensorFlow në operacione të bashkuara 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ë sfiduese për të identifikuar (p.sh. nëpërmjet përputhjes së modelit) nëngrafin që korrespondon me një operacion të tillë të përbërë.

  2. Mund të ketë më shumë se një zbatim TensorFlow që synon një operacion të bashkuar 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ë gjithë mund të konvertohen në të njëjtin operacion LSTM të shkrirë në LiteRT.

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

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

Në shumë raste, një pjesë e modelit mund të hartohet në një operacion të vetëm në TFLite. Kjo mund të ndihmojë me performancën kur shkruani një zbatim 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ë atribut tf.function me "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ë kaloni ato si pjesë e të njëjtave "implemente_eksperimentale".

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 nënkupton tashmë këtë.

Zbatoni opcionet e personalizuara dhe regjistrohuni me TFLite Interpreter

Zbatoni funksionimin tuaj të bashkuar si një funksion TFLite Custom - shikoni udhëzimet .

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

Një shembull për op 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 funksionimi i përbërë në funksion të shkrirë (i avancuar)

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

vizatim

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

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

Shkruani kodin e konvertimit

Kodi i konvertimit shkruhet në ndërfaqen e funksionit me shënimin implements . Shihni një shembull të bashkimit për kërkimin e ngulitjes . Konceptualisht, kodi i konvertimit zëvendëson zbatimin e përbërë të kësaj ndërfaqeje me atë të shkrirë.

Në kalimin përgatit-përbërës-funksionet, futni kodin tuaj të konvertimit .

Në përdorime më të avancuara, është e mundur të zbatohen transformime komplekse të operandeve të operacionit të përbërë për të nxjerrë operandët e operacionit të shkrirë. Shih Keras LSTM . kodi i konvertimit si shembull.

Konverto në LiteRT

Përdorni TFLiteConverter.from_saved_model API për ta kthyer në LiteRT.

Nën kapuç

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

Kompozimi i operacioneve në TensorFlow

Përdorimi i tf.function me atributin e funksionit eksperimental_implements i lejon përdoruesit të kompozojnë në mënyrë eksplicite operacione të reja duke përdorur operacionet primitive TensorFlow dhe të specifikojnë ndërfaqen që zbaton operacioni i përbërë rezultante. 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 ngulitjes. Kjo lidhet me një operacion të bashkuar 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ë bashkuara LiteRT.

Zgjerimi i konvertuesit LiteRT

Konvertuesi LiteRT që u lëshua në fillim të këtij viti mbështeti vetëm importimin e modeleve TensorFlow si grafik me të gjitha variablat të zëvendësuara me vlerat e tyre konstante përkatëse. Kjo nuk funksionon për shkrirjen e operacioneve pasi grafikë të tillë i kanë të gjitha funksionet të rreshtuara në mënyrë që variablat të mund të kthehen në konstante.

Në mënyrë që të shfrytëzohet funksioni 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 kemi zbatuar një rrjedhë të re pune të importimit dhe konvertimit të modeleve TensorFlow në konvertues për të mbështetur rastin e përdorimit të funksionit të kombinuar të shkrirjes. Konkretisht, veçoritë e reja të shtuara janë:

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

Kjo na lejon të kryejmë shkrirjen e operacioneve duke përdorur funksionet që përfaqësojnë operacionet e përbëra përpara inlinimit të funksionit dhe ngrirjes së ndryshueshme.

Zbatimi i funksionimit të bashkimit

Le të shohim më në detaje kalimin e funksionimit të shkrirjes. Ky kalim bën sa më poshtë:

  1. Hapni të gjitha funksionet në modulin MLIR.
  2. Nëse një funksion ka atributin tf._implements, bazuar në vlerën e atributit, thërret funksionin e duhur fusion utility.
  3. Programi i funksionimit të bashkimit operon në 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 shkrirë.
  4. Në shumë raste, trupi i zëvendësuar do të përmbajë operacione të tjera përveç funksionit 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ë palosen vazhdimisht, ato nuk do të ishin të pranishme në flatbuffer-in e eksportuar ku do të ekzistonte vetëm operacioni i bashkuar.

Këtu është një copë kodi nga leja 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 */
}

Këtu është një copë kodi që tregon hartimin e këtij operacioni të përbërë në një operacion të bashkuar në LiteRT duke shfrytëzuar 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());
  }