Halaman ini diterjemahkan oleh Cloud Translation API.
Switch to English

Menggunakan kompilasi AOT

Apa itu tfcompile?

tfcompile adalah alat mandiri yang sebelumnya (AOT) mengompilasi grafik TensorFlow menjadi kode yang dapat dieksekusi. Ini dapat mengurangi ukuran biner total, dan juga menghindari beberapa overhead waktu proses. Kasus penggunaan umum tfcompile adalah mengompilasi grafik inferensi menjadi kode yang dapat dijalankan untuk perangkat seluler.

Grafik TensorFlow biasanya dijalankan oleh waktu proses TensorFlow. Ini menimbulkan beberapa overhead waktu proses untuk eksekusi setiap node dalam grafik. Ini juga mengarah pada ukuran biner total yang lebih besar, karena kode untuk waktu proses 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.

Kompiler dibangun di atas kerangka XLA. Kode yang menghubungkan TensorFlow ke framework XLA berada di bawah tensorflow / compiler .

Apa yang dilakukan tfcompile?

tfcompile mengambil subgraf, yang diidentifikasi oleh konsep TensorFlow tentang umpan dan pengambilan, dan menghasilkan fungsi yang mengimplementasikan subgraf tersebut. The feeds adalah argumen input untuk fungsi, dan fetches adalah argumen output untuk fungsi. Semua masukan harus sepenuhnya ditentukan oleh feed; subgraf terpangkas yang dihasilkan tidak boleh berisi node Placeholder atau Variabel. Merupakan hal yang umum untuk menentukan semua Placeholder dan Variabel sebagai umpan, yang memastikan subgraf yang dihasilkan tidak lagi berisi node ini. Fungsi yang dihasilkan dikemas sebagai cc_library , dengan file header mengekspor tanda tangan fungsi, dan file objek yang berisi implementasi. Pengguna menulis kode untuk menjalankan fungsi yang dihasilkan sebagaimana mestinya.

Menggunakan tfcompile

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

  • Langkah 1: Konfigurasi subgraf yang akan dikompilasi
  • Langkah 2: Gunakan makro build tf_library untuk mengompilasi subgraf
  • Langkah 3: Tulis kode untuk menjalankan subgraf
  • Langkah 4: Buat biner terakhir

Langkah 1: Konfigurasi subgraf yang akan dikompilasi

Identifikasi umpan dan pengambilan yang sesuai dengan argumen masukan dan keluaran 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 makro build tf_library untuk mengompilasi subgraf

Langkah ini mengubah grafik menjadi cc_library menggunakan makro build 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 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 menghasilkan proto GraphDef (test_graph_tfmatmul.pb) untuk contoh ini, jalankan make_test_graphs.py dan tentukan lokasi keluaran 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. Demi kenyamanan, makro tf_library mendukung argumen freeze_checkpoint , yang menjalankan alat tersebut. Untuk contoh lainnya, lihat tensorflow / compiler / aot / tests / BUILD .

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

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

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

Langkah 3: Tulis kode untuk menjalankan subgraf

Langkah ini menggunakan file header ( test_graph_tfmatmul.h ) yang dibuat oleh makro build tf_library di langkah sebelumnya untuk memanggil kode yang dibuat. File header terletak di direktori bazel-bin sesuai dengan paket build, dan diberi nama berdasarkan atribut name yang disetel pada makro build tf_library . Misalnya, header yang dihasilkan untuk test_graph_tfmatmul akan menjadi 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 namespace foo::bar , karena itu adalah cpp_class ditentukan di makro tf_library . Semua kelas yang dihasilkan memiliki API serupa, dengan satu-satunya perbedaan adalah metode untuk menangani buffer arg 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 kelas yang dihasilkan: args merepresentasikan input, results merepresentasikan output, dan temps merepresentasikan buffer sementara yang digunakan secara internal untuk melakukan komputasi. 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 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 "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 pustaka 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",
    ]
)