Usando compilação AOT

Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

O que é tfcompile?

tfcompile é uma ferramenta autônoma que antes-do-time (AOT) compila TensorFlow gráficos em código executável. Ele pode reduzir o tamanho total do binário e também evitar algumas sobrecargas de tempo de execução. Um uso do caso típico de tfcompile é para compilar um gráfico inferência em código executável para dispositivos móveis.

O gráfico do TensorFlow normalmente é executado pelo tempo de execução do TensorFlow. Isso incorre em alguma sobrecarga de tempo de execução para a execução de cada nó no gráfico. Isso também leva a um tamanho binário total maior, uma vez que o código para o tempo de execução do TensorFlow precisa estar disponível, além do próprio gráfico. O código executável produzido por tfcompile não usa o tempo de execução TensorFlow, e só tem dependências em kernels que são efectivamente utilizadas no cálculo.

O compilador é construído com base na estrutura XLA. O código ponte TensorFlow aos reside quadro XLA sob tensorflow / compilador .

O que tfcompile faz?

tfcompile leva um subgrafo, identificado pelos conceitos TensorFlow de feeds e buscas, e gera uma função que implementos que subgráfico. Os feeds são os argumentos de entrada para a função, e as fetches são os argumentos de saída para a função. Todas as entradas devem ser totalmente especificadas pelos feeds; o subgrafo podado resultante não pode conter nós de espaço reservado ou variável. É comum especificar todos os marcadores de posição e variáveis ​​como feeds, o que garante que o subgráfico resultante não contenha mais esses nós. A função gerado é empacotado como um cc_library , com um ficheiro de cabeçalho exportar a assinatura de função, e um ficheiro de objecto contendo a aplicação. O usuário grava o código para invocar a função gerada conforme apropriado.

Usando tfcompile

Esta seção detalha as etapas de alto nível para a geração de um binário executável com tfcompile do subgrafo TensorFlow. As etapas são:

  • Etapa 1: configurar o subgráfico para compilar
  • Passo 2: Use o tf_library macro de construção para compilar o subgráfico
  • Etapa 3: escreva o código para invocar o subgráfico
  • Etapa 4: crie o binário final

Etapa 1: configurar o subgráfico para compilar

Identifique os feeds e buscas que correspondem aos argumentos de entrada e saída para a função gerada. Em seguida, configurar as feeds ea fetches em um tensorflow.tf2xla.Config proto.

# 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" }
}

Etapa 2: use a macro tf_library build para compilar o subgráfico

Esta etapa converte o gráfico em um cc_library usando o tf_library construção macro. O cc_library consiste em um ficheiro de objecto que contém o código gerado a partir do gráfico, juntamente com um ficheiro de cabeçalho que dá acesso ao código gerado. tf_library utiliza tfcompile para compilar o gráfico TensorFlow em código executável.

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

Para gerar o proto GraphDef (test_graph_tfmatmul.pb) para este exemplo, prazo make_test_graphs.py e especificar o local de saída com a bandeira --out_dir.

Gráficos típicos contêm Variables que representam os pesos que são aprendidas através da formação, mas tfcompile não é possível compilar um subgrafo que contêm Variables . O freeze_graph.py ferramenta converte variáveis em constantes, usando os valores armazenados em um arquivo de ponto de verificação. Como uma conveniência, o tf_library macro suporta o freeze_checkpoint argumento, que executa a ferramenta. Para mais exemplos ver tensorflow / compilador / aot / testes / construção .

As constantes que aparecem no subgráfico compilado são compiladas diretamente no código gerado. Para passar as constantes para a função gerada, em vez de compilá-las, simplesmente passe-as como feeds.

Para detalhes sobre o tf_library macro construção, ver tfcompile.bzl .

Para detalhes sobre o subjacente tfcompile ferramenta, consulte tfcompile_main.cc .

Etapa 3: escreva o código para invocar o subgráfico

Esta etapa usa o arquivo de cabeçalho ( test_graph_tfmatmul.h ) gerado pelo tf_library macro construção na etapa anterior para invocar o código gerado. O arquivo de cabeçalho está localizado no bazel-bin diretório correspondente ao pacote de construção, e é nomeado com base no conjunto nome do atributo na tf_library macro construção. Por exemplo, o cabeçalho gerado para test_graph_tfmatmul seria test_graph_tfmatmul.h . Abaixo está uma versão abreviada do que é gerado. O arquivo gerado, em bazel-bin , contém comentários adicionais úteis.

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

A classe gerada C ++ é chamado MatMulComp no foo::bar namespace, porque essa era a cpp_class especificado no tf_library macro. Todas as classes geradas têm uma API semelhante, com a única diferença sendo os métodos para lidar com arg e buffers de resultado. Esses métodos diferem de acordo com o número e tipos de tampões, que foram especificados pela feed e fetch argumentos para a tf_library macro.

Existem três tipos de buffers gerenciados dentro da classe gerada: args que representam os insumos, results que representam as saídas e temps representando buffers temporários usados internamente para executar o cálculo. Por padrão, cada instância da classe gerada aloca e gerencia todos esses buffers para você. O AllocMode argumento de construtor pode ser usado para alterar esse comportamento. Todos os buffers são alinhados aos limites de 64 bytes.

A classe C ++ gerada é apenas um invólucro em torno do código de baixo nível gerado pelo XLA.

Exemplo de invocar a função gerada com base em tfcompile_test.cc :

#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL

#include <iostream>
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "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;
}

Etapa 4: crie o binário final

Esta etapa combina a biblioteca gerada por tf_library no passo 2 e o código escrito na etapa 3 para criar um binário final. Abaixo está um exemplo bazel arquivo de construção.

# 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",
    ]
)