Questa pagina è stata tradotta dall'API Cloud Translation.

Utilizzo della compilazione AOT

Che cos'è tfcompile?

tfcompile è uno strumento autonomo che compila in anticipo i grafici TensorFlow in codice eseguibile. Può ridurre le dimensioni totali del programma binario ed evitare anche overhead del runtime. Un caso d'uso tipico di tfcompile è la compilazione di un grafico di inferenza nel codice eseguibile per i dispositivi mobili.

Il grafico TensorFlow viene solitamente eseguito dal runtime TensorFlow. Ciò comporta un overhead di runtime per l'esecuzione di ciascun nodo nel grafico. Ciò porta anche a una dimensione binaria totale maggiore, poiché deve essere disponibile il codice per il runtime di TensorFlow, oltre al grafico stesso. Il codice eseguibile prodotto da tfcompile non utilizza il runtime TensorFlow e ha dipendenze solo dai kernel effettivamente utilizzati nel calcolo.

Il compilatore è basato sul framework XLA. Il codice per il collegamento di TensorFlow al framework XLA si trova in tensorflow/compiler.

Che cosa fa tfcompile?

tfcompile prende un sottografico, identificato dai concetti di feed e recuperi di TensorFlow, e genera una funzione che implementa questo sottografico. feeds sono gli argomenti di input della funzione, mentre fetches sono gli argomenti di output per la funzione. Tutti gli input devono essere specificati completamente dai feed; il grafico secondario eliminato risultante non può contenere nodi segnaposto o variabili. È comune specificare tutti i segnaposto e le variabili come feed, in modo che il sottografico risultante non contenga più questi nodi. La funzione generata viene pacchettizzata come cc_library, con un file di intestazione che esporta la firma della funzione e un file di oggetto contenente l'implementazione. L'utente scrive il codice per richiamare la funzione generata in modo appropriato.

Utilizzo di tfcompile

Questa sezione descrive in dettaglio i passaggi di alto livello per generare un programma binario eseguibile con tfcompile da un sottografico di TensorFlow. I passaggi sono:

Passaggio 1: configura il sottografico per la compilazione
Passaggio 2: utilizza la macro di build tf_library per compilare il grafico secondario
Passaggio 3: scrivi il codice per richiamare il sottografico
Passaggio 4: crea il programma binario finale

Passaggio 1: configura il sottografico per la compilazione

Identifica i feed e i recuperi che corrispondono agli argomenti di input e output per la funzione generata. Quindi configura feeds e fetches in un protocollo tensorflow.tf2xla.Config.

# Each feed is a positional input argument for the generated function.  The order
# of each entry matches the order of each input argument.  Here “x_hold” and “y_hold”
# refer to the names of placeholder nodes defined in the graph.
feed {
  id { node_name: "x_hold" }
  shape {
    dim { size: 2 }
    dim { size: 3 }
  }
}
feed {
  id { node_name: "y_hold" }
  shape {
    dim { size: 3 }
    dim { size: 2 }
  }
}

# Each fetch is a positional output argument for the generated function.  The order
# of each entry matches the order of each output argument.  Here “x_y_prod”
# refers to the name of a matmul node defined in the graph.
fetch {
  id { node_name: "x_y_prod" }
}

Passaggio 2: utilizza la macro di compilazione tf_library per compilare il sottografo

Questo passaggio converte il grafico in un elemento cc_library utilizzando la macro di build tf_library. cc_library è costituito da un file oggetto contenente il codice generato dal grafico, insieme a un file di intestazione che consente di accedere al codice generato. tf_library utilizza tfcompile per compilare il grafico TensorFlow nel codice eseguibile.

load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# Use the tf_library macro to compile your graph into executable code.
tf_library(
    # name is used to generate the following underlying build rules:
    # <name>           : cc_library packaging the generated header and object files
    # <name>_test      : cc_test containing a simple test and benchmark
    # <name>_benchmark : cc_binary containing a stand-alone benchmark with minimal deps;
    #                    can be run on a mobile device
    name = "test_graph_tfmatmul",
    # cpp_class specifies the name of the generated C++ class, with namespaces allowed.
    # The class will be generated in the given namespace(s), or if no namespaces are
    # given, within the global namespace.
    cpp_class = "foo::bar::MatMulComp",
    # graph is the input GraphDef proto, by default expected in binary format.  To
    # use the text format instead, just use the ‘.pbtxt’ suffix.  A subgraph will be
    # created from this input graph, with feeds as inputs and fetches as outputs.
    # No Placeholder or Variable ops may exist in this subgraph.
    graph = "test_graph_tfmatmul.pb",
    # config is the input Config proto, by default expected in binary format.  To
    # use the text format instead, use the ‘.pbtxt’ suffix.  This is where the
    # feeds and fetches were specified above, in the previous step.
    config = "test_graph_tfmatmul.config.pbtxt",
)

Per generare il protocollo GraphDef (test_graph_tfmatmul.pb) per questo esempio, esegui make_test_graphs.py e specifica la posizione di output con il flag --out_dir.

I grafici tipici contengono Variables che rappresenta i pesi appresi tramite l'addestramento, ma tfcompile non può compilare un sottografico contenente Variables. Lo strumento freeze_graph.py converte le variabili in costanti, utilizzando i valori archiviati in un file di checkpoint. Per praticità, la macro tf_library supporta l'argomento freeze_checkpoint, che esegue lo strumento. Per altri esempi, consulta tensorflow/compiler/aot/tests/BUILD.

Le costanti visualizzate nel sottografico compilato vengono compilate direttamente nel codice generato. Per passare le costanti nella funzione generata, anziché completarle, trasmettile semplicemente come feed.

Per maggiori dettagli sulla macro di build tf_library, consulta tfcompile.bzl.

Per maggiori dettagli sullo strumento tfcompile sottostante, consulta tfcompile_main.cc.

Passaggio 3: scrivi il codice per richiamare il sottografico

Questo passaggio utilizza il file di intestazione (test_graph_tfmatmul.h) generato dalla macro di build tf_library nel passaggio precedente per richiamare il codice generato. Il file di intestazione si trova nella directory bazel-bin corrispondente al pacchetto di build e viene denominato in base all'attributo del nome impostato nella macro di compilazione tf_library. Ad esempio, l'intestazione generata per test_graph_tfmatmul sarebbe test_graph_tfmatmul.h. Di seguito è riportata una versione abbreviata di ciò che viene generato. Il file generato, in bazel-bin, contiene ulteriori commenti utili.

namespace foo {
namespace bar {

// MatMulComp represents a computation previously specified in a
// TensorFlow graph, now compiled into executable code.
class MatMulComp {
 public:
  // AllocMode controls the buffer allocation mode.
  enum class AllocMode {
    ARGS_RESULTS_AND_TEMPS,  // Allocate arg, result and temp buffers
    RESULTS_AND_TEMPS_ONLY,  // Only allocate result and temp buffers
  };

  MatMulComp(AllocMode mode = AllocMode::ARGS_RESULTS_AND_TEMPS);
  ~MatMulComp();

  // Runs the computation, with inputs read from arg buffers, and outputs
  // written to result buffers. Returns true on success and false on failure.
  bool Run();

  // Arg methods for managing input buffers. Buffers are in row-major order.
  // There is a set of methods for each positional argument.
  void** args();

  void set_arg0_data(float* data);
  float* arg0_data();
  float& arg0(size_t dim0, size_t dim1);

  void set_arg1_data(float* data);
  float* arg1_data();
  float& arg1(size_t dim0, size_t dim1);

  // Result methods for managing output buffers. Buffers are in row-major order.
  // Must only be called after a successful Run call. There is a set of methods
  // for each positional result.
  void** results();


  float* result0_data();
  float& result0(size_t dim0, size_t dim1);
};

}  // end namespace bar
}  // end namespace foo

La classe C++ generata è denominata MatMulComp nello spazio dei nomi foo::bar, perché era il valore cpp_class specificato nella macro tf_library. Tutte le classi generate hanno un'API simile. L'unica differenza sta nei metodi per gestire i buffer di argomento e risultato. Questi metodi differiscono in base al numero e ai tipi di buffer, che sono stati specificati dagli argomenti feed e fetch della macro tf_library.

Esistono tre tipi di buffer gestiti all'interno della classe generata: args che rappresenta gli input, results che rappresenta gli output e temps che rappresenta i buffer temporanei utilizzati internamente per eseguire il calcolo. Per impostazione predefinita, ogni istanza della classe generata alloca e gestisce tutti questi buffer per conto tuo. Puoi usare l'argomento del costruttore AllocMode per modificare questo comportamento. Tutti i buffer sono allineati ai limiti di 64 byte.

La classe C++ generata è solo un wrapper attorno al codice di basso livello generato da XLA.

Esempio di richiamo della funzione generata in base a tfcompile_test.cc:

#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL

#include <iostream>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "third_party/tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" // generated

int main(int argc, char** argv) {
  Eigen::ThreadPool tp(2);  // Size the thread pool as appropriate.
  Eigen::ThreadPoolDevice device(&tp, tp.NumThreads());


  foo::bar::MatMulComp matmul;
  matmul.set_thread_pool(&device);

  // Set up args and run the computation.
  const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
  std::copy(args + 0, args + 6, matmul.arg0_data());
  std::copy(args + 6, args + 12, matmul.arg1_data());
  matmul.Run();

  // Check result
  if (matmul.result0(0, 0) == 58) {
    std::cout << "Success" << std::endl;
  } else {
    std::cout << "Failed. Expected value 58 at 0,0. Got:"
              << matmul.result0(0, 0) << std::endl;
  }

  return 0;
}

Passaggio 4: crea il programma binario finale

Questo passaggio combina la libreria generata da tf_library nel passaggio 2 e il codice scritto nel passaggio 3 per creare il programma binario finale. Di seguito è riportato un esempio di file Build bazel.

# Example of linking your binary
# Also see //tensorflow/compiler/aot/tests/BUILD
load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library")

# The same tf_library call from step 2 above.
tf_library(
    name = "test_graph_tfmatmul",
    ...
)

# The executable code generated by tf_library can then be linked into your code.
cc_binary(
    name = "my_binary",
    srcs = [
        "my_code.cc",  # include test_graph_tfmatmul.h to access the generated header
    ],
    deps = [
        ":test_graph_tfmatmul",  # link in the generated object file
        "//third_party/eigen3",
    ],
    linkopts = [
          "-lpthread",
    ]
)