TensorFlowをC++でコーディングする方法

TensorFlowはC++で実装された機械学習フレームワークです。そのため、PythonだけでなくC++で実装することも可能です。C++で実装すれば、組み込みへの利用や、推論処理・学習処理の高速化といった恩恵があります。

一方で C++ で実装するにはいくつかハードルがあります。ハードルとしては、

1. 公式で C++ 用のライブラリがバイナリで提供されていないため、自分でビルドする必要がある
2. C++ API の公式ドキュメントでは、ほとんど使い方がわからない
3. インターネット上で公開されている情報が少ない

といったことが挙げられます。今回は C++ 実装に成功したので、その情報を書き記します。

ソースコードは shohata/tensorflow_perf で公開しています。

TensorFlow C++ライブラリのビルド

TensorFlowのC++用のライブラリは libtensorflow_cc.so と libtensorflow_framework.so の二つです。メインとして libtensorflow_cc.so があり、そのライブラリがコアである libtensorflow_framework.so に依存しています。

一般にインストールされるライブラリは libtensorflow.so であり、Python用のライブラリとして使われています。このライブラリをリンクしてC++ソースコードをコンパイルすることはほとんどできません（一部できるという記事も見つけたが再現性がありませんでした）。 libtensorflow_cc.so は残念ながら提供されていないので、自分でビルドする必要があります。

ビルドするには ソースからのビルド という公式ドキュメントが役立ちます。ただ、このドキュメントは libtensorflow.so をターゲットにしているため、そこを libtensorflow_cc.so に読み替える必要があります。

git clone -b v2.14.0 --depth 1 https://github.com/tensorflow/tensorflow.git
cd tensorflow
export TF_NEED_CUDA=1
export TF_CUDA_COMPUTE_CAPABILITIES="7.0,8.0"
./configure
bazel build --config=opt --config=cuda //tensorflow:libtensorflow_cc.so

このように bazel build の際にターゲットを //tensorflow:libtensorflow_cc.so とします。また、CUDAを使うように環境変数とオプションを指定します。この指定をしないと、CPU用にビルドされてしまい、GPUを利用することができません。 TF_CUDA_COMPUTE_CAPABILITIES でcapabilityを指定できるので、今回利用したいNVIDIA V100とNVIDIA A100が利用可能なように、 7.0,8.0 と指定します。この指定をしないと、JIT Compileが発生するため、起動に30分程度のコンパイル時間が発生してしまいます。

CUDAおよびcuDNNのバージョンは ソースからのビルド の下に表で記載されています。このバージョンを揃えないとコンパイルが通らないので、注意が必要です。

Version	Python version	Compiler	Build tools	cuDNN	CUDA
tensorflow-2.18.0	3.9-3.12	Clang 17.0.6	Bazel 6.5.0	9.3	12.5
tensorflow-2.17.0	3.9-3.12	Clang 17.0.6	Bazel 6.5.0	8.9	12.3
tensorflow-2.16.1	3.9-3.12	Clang 17.0.6	Bazel 6.5.0	8.9	12.3
tensorflow-2.15.0	3.9-3.11	Clang 16.0.0	Bazel 6.1.0	8.9	12.2
tensorflow-2.14.0	3.9-3.11	Clang 16.0.0	Bazel 6.1.0	8.7	11.8
tensorflow-2.13.0	3.8-3.11	Clang 16.0.0	Bazel 5.3.0	8.6	11.8

TensorFlow v2.14.0をターゲットとして nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 のDockerイメージでビルドしたところ、無事にビルドが通りました。その際に利用した Dockerfile がこちらです。

FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS clang-builder
ARG DEBIAN_FRONTEND=noninteractive
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        apt-transport-https \
        ca-certificates \
        gnupg \
        software-properties-common \
        wget \
    && rm -rf /var/lib/apt/lists/*
RUN wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
RUN add-apt-repository "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-16 main"
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        llvm-16 \
        clang-16 \
    && rm -rf /var/lib/apt/lists/*



FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS cmake-builder
ARG DEBIAN_FRONTEND=noninteractive
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        apt-transport-https \
        ca-certificates \
        gnupg \
        software-properties-common \
        wget \
    && rm -rf /var/lib/apt/lists/*
RUN wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | gpg --dearmor - | tee /etc/apt/trusted.gpg.d/kitware.gpg >/dev/null
RUN apt-add-repository "deb https://apt.kitware.com/ubuntu/ $(lsb_release -cs) main"
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        cmake \
    && rm -rf /var/lib/apt/lists/*



FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS bazel-builder
ARG DEBIAN_FRONTEND=noninteractive
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        apt-transport-https \
        ca-certificates \
        gnupg \
        curl \
    && rm -rf /var/lib/apt/lists/*
RUN curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor >bazel-archive-keyring.gpg
RUN mv bazel-archive-keyring.gpg /usr/share/keyrings
RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/bazel-archive-keyring.gpg] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        bazel-6.1.0 \
    && rm -rf /var/lib/apt/lists/*



FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS trt-builder
ARG DEBIAN_FRONTEND=noninteractive
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        apt-transport-https \
        ca-certificates \
        gnupg \
    && rm -rf /var/lib/apt/lists/*
ARG TRT_VERSION=8.6.1.6
ARG VER="${TRT_VERSION}-1+cuda11.8"
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub && \
    apt-get update && \
    apt-get install -y --no-install-recommends \
        libnvinfer8=${VER} libnvonnxparsers8=${VER} libnvparsers8=${VER} libnvinfer-plugin8=${VER} \
        libnvinfer-dev=${VER} libnvonnxparsers-dev=${VER} libnvparsers-dev=${VER} libnvinfer-plugin-dev=${VER} \
        python3-libnvinfer=${VER} libnvinfer-dispatch8=${VER} libnvinfer-dispatch-dev=${VER} libnvinfer-lean8=${VER} \
        libnvinfer-lean-dev=${VER} libnvinfer-vc-plugin8=${VER} libnvinfer-vc-plugin-dev=${VER} \
        libnvinfer-headers-dev=${VER} libnvinfer-headers-plugin-dev=${VER} \
    && rm -rf /var/lib/apt/lists/*



FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 AS tensorflow-builder
ARG DEBIAN_FRONTEND=noninteractive
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        ca-certificates \
        git \
        python3-dev \
        python3-pip \
        python-is-python3 \
    && rm -rf /var/lib/apt/lists/*

COPY --from=clang-builder /usr/lib/ /usr/lib/
COPY --from=clang-builder /usr/bin/ /usr/bin/
COPY --from=clang-builder /usr/bin/clang-16 /usr/bin/clang
COPY --from=clang-builder /usr/bin/clang++-16 /usr/bin/clang++
COPY --from=bazel-builder /usr/bin/bazel-6.1.0 /usr/bin/bazel
COPY --from=trt-builder /usr/lib/x86_64-linux-gnu/ /usr/lib/x86_64-linux-gnu/
COPY --from=trt-builder /lib/x86_64-linux-gnu/ /lib/x86_64-linux-gnu/
COPY --from=trt-builder /usr/include/x86_64-linux-gnu/ /usr/include/x86_64-linux-gnu/
RUN ldconfig

# Build TensorFlow
WORKDIR /
ARG TENSORFLOW_VERSION=v2.14.0
RUN git clone -b ${TENSORFLOW_VERSION} --depth 1 https://github.com/tensorflow/tensorflow.git
WORKDIR /tensorflow
ENV TF_NEED_CUDA=1
ENV TF_NEED_TENSORRT=1
ENV TF_CUDA_CLANG=1
ENV TF_NEED_ROCM=0
ENV TF_CUDA_COMPUTE_CAPABILITIES="7.0,8.0"
RUN ./configure
RUN bazel build --config=opt --config=cuda //tensorflow:libtensorflow_cc.so

# Copy TensorFlow C++ Library
ARG LIB_DIR=/usr/local/lib/
RUN cd bazel-bin/tensorflow && \
    ln -s libtensorflow_framework.so.2.14.0 libtensorflow_framework.so && \
    ln -s libtensorflow_framework.so.2.14.0 libtensorflow_framework.so.2 && \
    cp libtensorflow_cc.so ${LIB_DIR} && \
    cp libtensorflow_cc.so.2 ${LIB_DIR} && \
    cp libtensorflow_cc.so.2.14.0 ${LIB_DIR} && \
    cp libtensorflow_framework.so ${LIB_DIR} && \
    cp libtensorflow_framework.so.2 ${LIB_DIR} && \
    cp libtensorflow_framework.so.2.14.0 ${LIB_DIR}

# Copy TensorFlow C++ Header Files
ARG HEADER_DIR=/usr/local/include/
RUN mkdir -p ${HEADER_DIR}
RUN find tensorflow/core -follow -type f -name "*.h" -exec cp --parents {} ${HEADER_DIR} \;
RUN find tensorflow/cc -follow -type f -name "*.h" -exec cp --parents {} ${HEADER_DIR} \;
RUN find tensorflow/c -follow -type f -name "*.h" -exec cp --parents {} ${HEADER_DIR} \;
RUN find tensorflow/tsl -follow -type f -name "*.h" -exec cp --parents {} ${HEADER_DIR} \;
RUN find third_party/eigen3 -follow -type f -exec cp --parents {} ${HEADER_DIR} \;
RUN cd bazel-bin && \
    find tensorflow -follow -type f -name "*.h" -exec cp --parents {} ${HEADER_DIR} \;
RUN cd bazel-tensorflow/external/com_google_protobuf/src && \
    find google -follow -type f -name "*.h" -exec cp --parents {} ${HEADER_DIR} \; && \
    find google -follow -type f -name "*.inc" -exec cp --parents {} ${HEADER_DIR} \;
RUN cd bazel-tensorflow/external/com_google_absl && \
    find absl -follow -type f -exec cp --parents {} ${HEADER_DIR} \;
RUN cd bazel-tensorflow/external/ml_dtypes && \
    find include -follow -type f -exec cp --parents {} ${HEADER_DIR} \;
RUN cd bazel-tensorflow/external/eigen_archive && \
    find Eigen -follow -type f -exec cp --parents {} ${HEADER_DIR} \; && \
    find unsupported -follow -type f -exec cp --parents {} ${HEADER_DIR} \;


FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04
ARG DEBIAN_FRONTEND=noninteractive
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        build-essential \
        ninja-build \
        ca-certificates \
        git \
        python3-dev \
        python3-pip \
        python-is-python3 \
    && rm -rf /var/lib/apt/lists/*

ARG CMAKE_VERSION=3.27
COPY --from=clang-builder /usr/lib/ /usr/lib/
COPY --from=clang-builder /usr/bin/ /usr/bin/
COPY --from=clang-builder /usr/bin/clang-16 /usr/bin/clang
COPY --from=clang-builder /usr/bin/clang++-16 /usr/bin/clang++
COPY --from=cmake-builder /usr/bin/cmake /usr/bin/cmake
COPY --from=cmake-builder /usr/share/cmake-${CMAKE_VERSION}/ /usr/share/cmake-${CMAKE_VERSION}/
COPY --from=trt-builder /usr/lib/x86_64-linux-gnu/ /usr/lib/x86_64-linux-gnu/
COPY --from=trt-builder /lib/x86_64-linux-gnu/ /lib/x86_64-linux-lib
COPY --from=trt-builder /usr/include/x86_64-linux-gnu/ /usr/include/x86_64-linux-gnu/
COPY --from=tensorflow-builder /usr/local/lib/ /usr/local/lib/
COPY --from=tensorflow-builder /usr/local/include/ /usr/local/include/
RUN ldconfig

COPY FindTensorFlow.cmake /usr/share/cmake-${CMAKE_VERSION}/Modules/

RUN pip install --upgrade pip
RUN pip install tensorflow
RUN pip install --upgrade tensorflow-hub

ENV CC=/usr/bin/clang
ENV CXX=/usr/bin/clang++

TensorFlow C++ API

公式のドキュメントして TensorFlow C++ API Reference が公開されています。使い方はGitHubの tensorflow/tensorflow リポジトリにある tensorflow/examples/label_image/main.cc が参考になります。画像を読み込んで、処理する例を挙げます。

using namespace tensorflow::ops;

tensorflow::Scope root = tensorflow::Scope::NewRootScope();

const std::string input_name = "filename";
const std::string output_name = "processed_image";
const int32_t input_height = 224;
const int32_t input_width = 224;
const int64_t input_channels = 3;

auto filename = Placeholder(root.WithOpName(input_name), tensorflow::DT_STRING);
auto jpeg_image = ReadFile(root.WithOpName("read_file"), filename);
auto decode_image = DecodeJpeg(root.WithOpName("decode_jpeg"), jpeg_image, DecodeJpeg::Attrs().Channels(input_channels));
auto expand_image = ExpandDims(root.WithOpName("expand_dimension"), decode_image, 0);
auto resize_image = ResizeBilinear(root.WithOpName("resize_image"), expand_image, GuaranteeConst(root.WithOpName("size"), {input_height, input_width}));
auto cast_image = Cast(root.WithOpName(output_name), resize_image, tensorflow::DT_FLOAT);

// Create Image Processing Graph
tensorflow::GraphDef graph;
tensorflow::Session *image_session = tensorflow::NewSession(tensorflow::SessionOptions());
TF_RETURN_IF_ERROR(root.ToGraphDef(&graph));
TF_RETURN_IF_ERROR(session->Create(graph));

const std::string image_path = "/path/to/image.jpg";
std::vector<Tensor> outputs;

TF_RETURN_IF_ERROR(session->Run({{input_name, Tensor(image_path)}}, {output_name}, {}, &outputs));

このコードはグラフ（機械学習モデル）を作成し、実行（推論）するコードと全く同じです。作成されるグラフ（MLモデル）がJPEG画像ファイル名をインプット、前処理された行列（Tensor）をアウトプット、とするグラフというわけです。そのグラフに対して推論処理をすることで、推論結果に代わり前処理結果が受け取れます。

公式のドキュメント TensorFlow C++ API Reference で公開されているAPIはグラフを構成するオペレーションになります。通常の機械学習モデルに使用されるDNNオペレーションから、ファイル読み込みや画像処理まで、各種オペレーションを備えています。このオペレーションを組み合わせることで、推論から前処理まで各種操作が可能となります。

Saved Modelを使った推論

TensorFlowの学習済み機械学習モデルの保存形式としてSaved Modelがあります。公式ドキュメントではSavedModel形式の使用例としてPythonでの使用方法とC++での読み込み方法が掲載されています。ただC++での読み込みの情報が少なすぎてこれだけでは使えないので、完全なコードを掲載します。

tensorflow::SavedModelBundle bundle;
const std::string model_path ="/saved_model/resnet_50"

TF_RETURN_IF_ERROR(
      tensorflow::LoadSavedModel(
          tensorflow::SessionOptions(), tensorflow::RunOptions(), model_path,
          {tensorflow::kSavedModelTagServe}, &bundle));

tensorflow::Session *session = bundle.GetSession();
auto model_def = bundle.GetSignatures().at("serving_default");
const std::string input_name = model_def.inputs().at("keras_layer_input").name();
const std::string output_name = model_def.outputs().at("keras_layer").name();

// A image which get from previous code `outputs`
Tensor image;
std::vector<Tensor> outputs;

TF_RETURN_IF_ERROR(session->Run({{input_name, image}}, {output_name}, {}, &outputs));

tensorflow::LoadSavedModel を利用して bundle に学習済みモデルを読み込みます。 bundle.GetSession() を利用してセッションを取得すれば、後は同様に session->Run() で推論することができます。入力と出力のオペレーション名は、保存されたSaved Modelの入出力オペレーション名と揃える必要があります。このコードではデフォルトの名前を使用しています。

まとめ

TensorFlow C++開発環境構築から画像処理、推論処理に至るまで、一連の流れをまとめました。このコードを使うことで、Pythonで保存したSaved Modelに対して推論処理をすることができます。高速に推論処理を行いたい場合に有効活用するのはいかがでしょうか。