Halaman ini diterjemahkan oleh Cloud Translation API.
Switch to English

Menggunakan kompilasi AOT

Apa itu tfcompile?

tfcompile adalah alat mandiri yang sebelumnya (AOT) mengkompilasi grafik TensorFlow menjadi kode yang dapat dieksekusi. Ini dapat mengurangi ukuran biner total, dan juga menghindari beberapa overhead runtime. Contoh penggunaan tfcompile adalah mengkompilasi grafik inferensi ke dalam kode yang dapat dieksekusi untuk perangkat seluler.

Grafik TensorFlow biasanya dijalankan oleh runtime TensorFlow. Ini menimbulkan beberapa overhead runtime untuk mengeksekusi setiap node dalam grafik. Ini juga mengarah ke ukuran biner total yang lebih besar, karena kode untuk runtime TensorFlow harus tersedia, selain grafik itu sendiri. Kode yang dapat dieksekusi yang dihasilkan oleh tfcompile tidak menggunakan runtime TensorFlow, dan hanya memiliki dependensi pada kernel yang benar-benar digunakan dalam perhitungan.

Kompiler dibuat di atas kerangka XLA. Kode menjembatani TensorFlow ke kerangka XLA berada di bawah tensorflow / compiler .

Apa yang dilakukan tfcompile?

tfcompile mengambil subgraph, diidentifikasi oleh konsep TensorFlow dari feed dan fetch, dan menghasilkan fungsi yang mengimplementasikan subgraph itu. The feeds adalah argumen input untuk fungsi, dan fetches adalah argumen output untuk fungsi. Semua input harus ditentukan sepenuhnya oleh feed; subgraph prun yang dihasilkan tidak dapat berisi node Placeholder atau Variabel. Merupakan hal yang umum untuk menentukan semua Placeholder dan Variabel sebagai feed, yang memastikan subgraph yang dihasilkan tidak lagi berisi node ini. Fungsi yang dihasilkan dikemas sebagai cc_library , dengan file header yang mengekspor tanda tangan fungsi, dan file objek yang berisi implementasinya. Pengguna menulis kode untuk menjalankan fungsi yang dihasilkan sebagaimana mestinya.

Menggunakan tfcompile

Bagian ini merinci langkah-langkah tingkat tinggi untuk menghasilkan biner yang dapat dieksekusi dengan tfcompile dari subgraph TensorFlow. Langkah-langkahnya adalah:

  • Langkah 1: Konfigurasikan subgraph untuk dikompilasi
  • Langkah 2: Gunakan tf_library build macro untuk mengkompilasi subgraph
  • Langkah 3: Tulis kode untuk menjalankan subgraph
  • Langkah 4: Buat biner terakhir

Langkah 1: Konfigurasikan subgraph untuk dikompilasi

Identifikasi umpan dan pengambilan yang sesuai dengan argumen input dan output untuk fungsi yang dihasilkan. Kemudian mengkonfigurasi feeds dan fetches dalam 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" }
}
 

Langkah 2: Gunakan tf_library build macro untuk mengkompilasi subgraph

Langkah ini mengubah grafik menjadi cc_library menggunakan makro tf_library . cc_library terdiri dari file objek yang berisi kode yang dihasilkan dari grafik, bersama dengan file header yang memberikan akses ke kode yang dihasilkan. tf_library menggunakan tfcompile untuk mengkompilasi grafik TensorFlow menjadi kode yang dapat dieksekusi.

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

Untuk menghasilkan proto GraphDef (test_graph_tfmatmul.pb) untuk contoh ini, jalankan make_test_graphs.py dan tentukan lokasi output dengan flag --out_dir.

Grafik tipikal berisi Variables mewakili bobot yang dipelajari melalui pelatihan, tetapi tfcompile tidak dapat mengkompilasi subgraf yang berisi Variables . Alat freeze_graph.py mengubah variabel menjadi konstanta, menggunakan nilai yang disimpan dalam file pos pemeriksaan. Sebagai kenyamanan, makro tf_library mendukung argumen freeze_checkpoint , yang menjalankan alat. Untuk lebih banyak contoh, lihat tensorflow / compiler / aot / tests / BUILD .

Konstanta yang muncul dalam subgraph yang dikompilasi dikompilasi langsung ke dalam kode yang dihasilkan. Untuk meneruskan konstanta ke dalam fungsi yang dihasilkan, alih-alih meminta mereka dikompilasi, cukup kirimkan sebagai umpan.

Untuk detail tentang makro tf_library build, lihat tfcompile.bzl .

Untuk detail tentang alat tfcompile mendasarinya, lihat tfcompile_main.cc .

Langkah 3: Tulis kode untuk menjalankan subgraph

Langkah ini menggunakan file header ( test_graph_tfmatmul.h ) yang dihasilkan oleh tf_library build macro pada langkah sebelumnya untuk memanggil kode yang dihasilkan. File header terletak di direktori bazel-bin sesuai dengan paket build, dan dinamai berdasarkan atribut nama yang ditetapkan pada makro build tf_library . Misalnya, tajuk yang dibuat untuk test_graph_tfmatmul adalah test_graph_tfmatmul.h . Di bawah ini adalah versi singkat dari apa yang dihasilkan. File yang dihasilkan, di bazel-bin , berisi komentar berguna tambahan.

 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
 

Kelas C ++ yang dihasilkan disebut MatMulComp di foo::bar namespace, karena itu adalah cpp_class ditentukan dalam makro tf_library . Semua kelas yang dihasilkan memiliki API yang sama, dengan satu-satunya perbedaan adalah metode untuk menangani arg dan menghasilkan buffer. Metode-metode itu berbeda berdasarkan jumlah dan jenis buffer, yang ditentukan oleh feed dan fetch argumen ke makro tf_library .

Ada tiga jenis buffer yang dikelola dalam kelas yang dihasilkan: args mewakili input, results mewakili output, dan temps mewakili buffer sementara yang digunakan secara internal untuk melakukan perhitungan. Secara default, setiap instance dari kelas yang dihasilkan mengalokasikan dan mengelola semua buffer ini untuk Anda. AllocMode konstruktor AllocMode dapat digunakan untuk mengubah perilaku ini. Semua buffer disejajarkan dengan batas 64-byte.

Kelas C ++ yang dihasilkan hanyalah pembungkus di sekitar kode tingkat rendah yang dihasilkan oleh XLA.

Contoh memanggil fungsi yang dihasilkan berdasarkan 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;
}
 

Langkah 4: Buat biner terakhir

Langkah ini menggabungkan perpustakaan yang dihasilkan oleh tf_library di langkah 2 dan kode yang ditulis di langkah 3 untuk membuat biner akhir. Di bawah ini adalah contoh 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",
    ]
)