pybind11 + pytorch + cmake で pytorch(Python) + C++ 連携のメモ #PyTorch

背景

pytorch 1.8 から fft(CUDA 対応) やら linalg(beta version) ついたりして, pytorch を数値演算フレームワークとして使うでもいい感じなので, python + C++ アプリに pytorch を pybind11 で組み込みたい.

Python 世界だけだと, 大規模データや並列処理がやはりどうしても苦手になってしまう. また Multiprocessing とか Pipe だと無駄にメモリを消費したり扱えるデータサイズに制限があったりでやっぱり大規模には向かない.

必要なところだけ C++ で効率的にやりつつも, 残りは pytorch つかって効率的にコーディングしたい.

Ubuntu(Linux)を想定します.

conda installed pytorch

python 環境との親和性を考えて, libtorch ではなく conda や pip で入る通常の pytorch を利用します.

conda などで python 環境を作っているとします.

TorchConfig.cmake の探索

以下のように python site パッケージで pytorch のインストールしてあるパスを取得して指定します.

cmake で Torch_DIR で

TORCH_INSTALLED_DIR=`python -c "import site; print (site.getsitepackages()[0])"`/torch/share/cmake

cmake -DTorch_DIR=${TORCH_INSTALLED_DIR} ...

cmake では, Xyz_DIR とすると, XyzConfig.cmake or xyz-config.cmake を探すようになっています.

CMAKE_MODULE_PATH で指定もできますが, 他のカスタム cmake ファイルの探索とかぶってうまくいかないときもありますから, Torch_DIR がよいでしょう.

Torch_DIR は関連する cmake オプションより手前に定義する必要がありますので, 環境変数で処理するのがいいかもしれません.
(CMAKE_MODULE_PATH は優先順位高いのか cmake 内で他のオプションより先に解釈されるっぽい)

Torch_DIR=${TORCH_INSTALLED_DIR} cmake ....

C++ compiler ABI

conda とかで入る pytorch は互換性のためか new GCC CXX11 ABI off(昔の ABI)でコンパイルされています. コンパイラも gcc を使っているようです.

# When we build libtorch with the old GCC ABI, dependent libraries must too.
if("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU")
  set(TORCH_CXX_FLAGS "-D_GLIBCXX_USE_CXX11_ABI=0")
endif()

TorchConfig.cmake では, gcc の場合は自動で _GLIBCXX_USE_CXX11_ABI=0 を define しますが, clang などの場合はこのフラグを明示的に target_compile_options で追加する必要があります.

定義しないと, コンパイルはいけますが実行時にクラッシュします.

C++ STL

また, STL は gnustl なので, libc++ を使いたい場合は pytorch をソースコードからビルドする必要があるでしょう.

リンク

TorchConfig.cmake では, torch_python とのリンクが抜けていますので追加します.

lib のパスは

${TORCH_INSTALL_PREFIX}/lib で得られます.

例

最終的に CMake スクリプトは以下のようになります.

find_package(Torch REQUIRED)

# path to pybind11
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third_party/pybind11 pybind11_build)

add_library(torch_testmod SHARED
    src/torch_library.cc)

target_include_directories(torch_testmod PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/pybind11/include)

target_compile_options(torch_testmod PRIVATE "-D_GLIBCXX_USE_CXX11_ABI=0")

target_link_libraries(torch_testmod "${TORCH_LIBRARIES}"
    torch_python
    pybind11::module pybind11::lto pybind11::windows_extras)
set_property(TARGET torch_testmod PROPERTY CXX_STANDARD 14)

# TORCH_INSTALL_PREFIX will be available in find_package(Torch)
target_link_directories(torch_testmod PRIVATE ${TORCH_INSTALL_PREFIX}/lib)

pybind11_extension(torch_testmod)
pybind11_strip(torch_testmod)

C++ コード例. torch のヘッダのどこかに pybind11 インクルードがあるようで, pybind11 ヘッダの明示的なインクルードは不要

#include <torch/extension.h>

#include <iostream>

torch::Tensor d_sigmoid(torch::Tensor z) {
  auto s = torch::sigmoid(z);
  return (1 - s) * s;
}

PYBIND11_MODULE(torch_testmod, m) {
  m.def("sigmoid", &d_sigmoid, "Sigmoid");
}

cmake bootstrap は以下のようになります.


# Path to TorchConfig cmake files.
TORCH_INSTALLED_DIR=`python -c "import site; print (site.getsitepackages()[0])"`/torch/share/cmake

echo "TORCH_DIR = " ${TORCH_INSTALLED_DIR}

mkdir build
cd build

Torch_DIR=${TORCH_INSTALLED_DIR} \
CXX=clang++-11 CC=clang-11 cmake -DCMAKE_BUILD_TYPE=Debug \
  -DCMAKE_VERBOSE_MAKEFILE=1 \
  ..

ビルドしたあとの, サンプル .py は以下のようになります.

>>> import torch
>>> import torch_testmod
>>> a = torch.Tensor([1, 2, 3])
>>> torch_testmod.sigmoid(a)
tensor([0.1966, 0.1050, 0.0452])

Voila!

問題点

libtorch および pybind11 がヘッダたくさん, テンプレートたくさんでコンパイルが激遅です.

上記の数行の C++ コードでも Threadripper 1950X で 10 秒くらいかかります.

torch と関連するコードはなるべく変更なしにして, ccache でコンパイルキャッシュ作って速くするなどの対応が求められるでしょう.