이 페이지는 Cloud Translation API를 통해 번역되었습니다.

AOT 컴파일 사용

tfcompile이란?

tfcompile는 AOT (ahead-of-time) 기능이 TensorFlow 그래프를 실행 코드로 컴파일하는 독립형 도구입니다. 총 바이너리 크기를 줄이고 일부 런타임 오버헤드를 피할 수 있습니다. tfcompile의 일반적인 사용 사례는 추론 그래프를 휴대기기의 실행 코드로 컴파일하는 것입니다.

TensorFlow 그래프는 일반적으로 TensorFlow 런타임에 의해 실행됩니다. 이로 인해 그래프의 각 노드를 실행하는 데 런타임 오버헤드가 발생합니다. 또한 그래프 자체 외에도 TensorFlow 런타임용 코드를 사용할 수 있어야 하므로 총 바이너리 크기가 더 커집니다. tfcompile에 의해 생성된 실행 코드는 TensorFlow 런타임을 사용하지 않으며 계산에 실제로 사용되는 커널에 관한 종속 항목만 있습니다.

컴파일러는 XLA 프레임워크를 기반으로 빌드됩니다. TensorFlow를 XLA 프레임워크에 연결하는 코드는 tensorflow/compiler에 있습니다.

tfcompile로 무엇을 수행하나요?

tfcompile는 TensorFlow의 피드 및 가져오기 개념으로 식별되는 하위 그래프를 가져와서 해당 하위 그래프를 구현하는 함수를 생성합니다. feeds는 함수의 입력 인수이고 fetches는 함수의 출력 인수입니다. 모든 입력은 피드에서 완전히 지정해야 합니다. 결과적으로 프루닝된 하위 그래프에는 자리표시자 또는 변수 노드가 포함될 수 없습니다. 일반적으로 모든 자리표시자와 변수를 피드로 지정하여 결과 하위 그래프에 더 이상 이러한 노드가 포함되지 않도록 합니다. 생성된 함수는 함수 서명을 내보내는 헤더 파일과 구현을 포함하는 객체 파일과 함께 cc_library로 패키징됩니다. 사용자는 코드를 작성하여 생성된 함수를 적절하게 호출합니다.

tfcompile 사용

이 섹션에서는 TensorFlow 하위 그래프에서 tfcompile를 사용하여 실행 바이너리를 생성하는 대략적인 단계를 자세히 설명합니다. 단계는 다음과 같습니다.

1단계: 하위 그래프 컴파일하기
2단계: tf_library 빌드 매크로를 사용하여 하위 그래프 컴파일하기
3단계: 하위 그래프를 호출하는 코드 작성
4단계: 최종 바이너리 만들기

1단계: 하위 그래프 컴파일하기

생성된 함수의 입력 및 출력 인수에 해당하는 피드와 가져오기를 식별합니다. 그런 다음 tensorflow.tf2xla.Config proto에서 feeds 및 fetches를 구성합니다.

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

2단계: tf_library 빌드 매크로를 사용하여 하위 그래프 컴파일

이 단계에서는 tf_library 빌드 매크로를 사용하여 그래프를 cc_library로 변환합니다. cc_library는 그래프에서 생성된 코드가 포함된 객체 파일과 생성된 코드에 대한 액세스 권한을 부여하는 헤더 파일로 구성됩니다. tf_library는 tfcompile를 활용하여 TensorFlow 그래프를 실행 코드로 컴파일합니다.

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

이 예시의 GraphDef proto (test_graph_tfmatmul.pb)를 생성하려면 make_test_graphs.py를 실행하고 --out_dir 플래그를 사용하여 출력 위치를 지정합니다.

일반적인 그래프에는 학습을 통해 학습된 가중치를 나타내는 Variables가 포함되지만 tfcompile는 Variables가 포함된 하위 그래프를 컴파일할 수 없습니다. freeze_graph.py 도구는 체크포인트 파일에 저장된 값을 사용하여 변수를 상수로 변환합니다. 편의를 위해 tf_library 매크로는 도구를 실행하는 freeze_checkpoint 인수를 지원합니다. 더 많은 예는 tensorflow/compiler/aot/tests/BUILD를 참조하세요.

컴파일된 하위 그래프에 표시되는 상수는 생성된 코드로 직접 컴파일됩니다. 상수를 컴파일된 함수에 전달하는 대신 생성된 함수에 전달하려면 피드로 전달하기만 하면 됩니다.

tf_library 빌드 매크로에 관한 자세한 내용은 tfcompile.bzl을 참조하세요.

기본 tfcompile 도구에 관한 자세한 내용은 tfcompile_main.cc를 참조하세요.

3단계: 하위 그래프를 호출하는 코드 작성

이 단계에서는 이전 단계의 tf_library 빌드 매크로로 생성된 헤더 파일 (test_graph_tfmatmul.h)을 사용하여 생성된 코드를 호출합니다. 헤더 파일은 빌드 패키지에 해당하는 bazel-bin 디렉터리에 있으며 tf_library 빌드 매크로에 설정된 이름 속성을 기반으로 이름이 지정됩니다. 예를 들어 test_graph_tfmatmul용으로 생성된 헤더는 test_graph_tfmatmul.h입니다. 아래는 생성되는 결과의 축약된 버전입니다. 생성된 파일(bazel-bin)에는 유용한 추가 주석이 포함되어 있습니다.

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

생성된 C++ 클래스는 tf_library 매크로에 지정된 cpp_class였으므로 foo::bar 네임스페이스에서 MatMulComp라고 합니다. 생성된 모든 클래스에는 비슷한 API가 있지만 인수 및 결과 버퍼를 처리하는 메서드만 다릅니다. 이러한 메서드는 tf_library 매크로의 feed 및 fetch 인수로 지정된 버퍼의 수와 유형에 따라 다릅니다.

생성된 클래스 내에서 관리되는 버퍼에는 세 가지 유형이 있습니다. 입력을 나타내는 args, 출력을 나타내는 results, 계산 실행을 위해 내부적으로 사용되는 임시 버퍼를 나타내는 temps입니다. 기본적으로 생성된 클래스의 각 인스턴스는 이러한 모든 버퍼를 할당하고 관리합니다. AllocMode 생성자 인수를 사용하여 이 동작을 변경할 수 있습니다. 모든 버퍼는 64바이트 경계에 맞춰 정렬됩니다.

생성된 C++ 클래스는 XLA에 의해 생성된 하위 수준 코드를 둘러싸는 래퍼일 뿐입니다.

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;
}

4단계: 최종 바이너리 만들기

이 단계에서는 2단계에서 tf_library에 의해 생성된 라이브러리와 3단계에서 작성된 코드를 결합하여 최종 바이너리를 만듭니다. 다음은 bazel BUILD 파일의 예입니다.

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