Halaman ini diterjemahkan oleh Cloud Translation API.

Menggunakan kompilasi AOT

Apa itu tfcompile?

tfcompile adalah alat mandiri yang ahead-of-time (AOT) mengompilasi grafik TensorFlow menjadi kode yang dapat dieksekusi. Hal ini dapat mengurangi total ukuran biner, dan juga menghindari beberapa overhead runtime. Kasus penggunaan standar tfcompile adalah mengompilasi grafik inferensi menjadi kode yang dapat dieksekusi untuk perangkat seluler.

Grafik TensorFlow biasanya dijalankan oleh runtime TensorFlow. Hal ini menimbulkan beberapa overhead runtime untuk eksekusi setiap node dalam grafik. Hal ini juga menyebabkan total ukuran biner 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 komputasi.

Compiler dibuat di atas framework XLA. Kode yang menjembatani TensorFlow ke framework XLA berada di tensorflow/compiler.

Apa yang dilakukan tfcompile?

tfcompile mengambil subgrafik, yang diidentifikasi oleh konsep TensorFlow terkait feed dan pengambilan, lalu menghasilkan fungsi yang menerapkan subgrafik tersebut. feeds adalah argumen input untuk fungsi, dan fetches adalah argumen output untuk fungsi. Semua input harus ditentukan sepenuhnya oleh feed; subgrafik yang dipangkas yang dihasilkan tidak boleh berisi node Placeholder atau Variabel. Menentukan semua Placeholder dan Variabel sebagai feed adalah hal yang umum dilakukan, yang memastikan subgrafik 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 implementasi. Pengguna menulis kode untuk memanggil fungsi yang dihasilkan sebagaimana mestinya.

Menggunakan tfcompile

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

Langkah 1: Konfigurasi subgrafik untuk dikompilasi
Langkah 2: Gunakan makro build tf_library untuk mengompilasi subgrafik
Langkah 3: Tulis kode untuk memanggil subgrafik
Langkah 4: Buat biner akhir

Langkah 1: Konfigurasi subgrafik untuk dikompilasi

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

Langkah 2: Gunakan makro build tf_library untuk mengompilasi subgrafik

Langkah ini mengonversi grafik menjadi cc_library menggunakan makro build tf_library. cc_library terdiri dari file objek yang berisi kode yang dihasilkan dari grafik, beserta file header yang memberikan akses ke kode yang dihasilkan. tf_library menggunakan tfcompile untuk mengompilasi 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 membuat proto GraphDef (test_graph_tfmatmul.pb) untuk contoh ini, jalankan make_test_graphs.py dan tentukan lokasi output dengan tanda --out_dir.

Grafik standar berisi Variables yang mewakili bobot yang dipelajari melalui pelatihan, tetapi tfcompile tidak dapat mengompilasi subgrafik yang berisi Variables. Alat freeze_graph.py mengonversi variabel menjadi konstanta, menggunakan nilai yang disimpan dalam file checkpoint. Untuk memudahkan, makro tf_library mendukung argumen freeze_checkpoint, yang menjalankan alat tersebut. Untuk contoh lainnya, lihat tensorflow/compiler/aot/tests/BUILD.

Konstanta yang muncul dalam subgrafik yang dikompilasi dikompilasi langsung ke dalam kode yang dihasilkan. Untuk meneruskan konstanta ke dalam fungsi yang dihasilkan, daripada mengompilasinya, cukup teruskan konstanta sebagai feed.

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

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

Langkah 3: Tulis kode untuk memanggil subgrafik

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

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

Class C++ yang dihasilkan disebut MatMulComp dalam namespace foo::bar, karena itu adalah cpp_class yang ditentukan dalam makro tf_library. Semua class yang dihasilkan memiliki API serupa, dengan satu-satunya perbedaan adalah metode untuk menangani buffer argumen dan hasil. Metode tersebut berbeda berdasarkan jumlah dan jenis buffer, yang ditentukan oleh argumen feed dan fetch ke makro tf_library.

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

Class C++ yang dihasilkan hanyalah wrapper di sekitar kode tingkat rendah yang dihasilkan oleh XDLA.

Contoh pemanggilan 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 "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;
}

Langkah 4: Buat biner akhir

Langkah ini menggabungkan library yang dihasilkan oleh tf_library pada langkah 2 dan kode yang ditulis pada 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",
    ]
)