次世代Pythonバインディングライブラリnanobindを試す

Pythonバインディングとは？

(主に)C/C++で書かれたコードをpythonで使えるようにすることを言います。
このようなことを行う理由としては、Pythonはもともと計算のパフォーマンスがあまり良くなく、ボトルネックとなるような処理をC/C++のような処理の速いコンパイル言語で記述し、それをPythonから呼び出せるようにするためで、これによって計算のボトルネックを解消することがしばしば行われます。
C/C++以外でもバインディングすることもあるかと思いますが、C/C++だとSIMDやOpenMP、CUDAといったハードウェアアクセラレーションとの親和性も高く、特に科学技術系ではこのような高速化が行われることが多いです。

Pythonバインディングを行うためのライブラリ

本記事で取り上げるnanobindを紹介する前に、現在よく使用されている他のPythonバインディングライブラリを紹介します。

Pybind11

世の中のライブラリを見ていると、これか次のCythonを使っているものが多いのではないかという印象です。有名所でいくとPytorchやTensorflowがpybind11を採用しています。
バインディング部分含めてすべてC++11で記述することができ、poetryやsetup.pyを用いたビルドまわりのサポートもされています。

Cython

Pybind11が比較的新しめのライブラリなので、その前からあるライブラリではCythonを使用しているものが多いイメージです。有名所でいくとNumPyやCuPy、NVIDIAのRAPIDSプロジェクトが採用しています。
Pythonに型情報を付加したような独自の言語で記述することで、C/C++のバインディングを行うことができます。C/C++に慣れない人には使いやすいと思われます。

ctypes

C++を使わない純粋なCのライブラリであれば、Pythonの標準ライブラリに用意されているctypesモジュールを使ってPythonから直接ダイナミックリンクライブラリを読み込むことができます。

PyO3

こちらはC/C++ではなく、RustのPythonバインディングライブラリです。
maturinというツールを使うことで、プロジェクトの作成やCIの設定まわりを自動で行ってくれます。
Rustは最近人気が出てきいる言語なので、今後、Pythonバインディングを書く際の選択肢として入ってくることは多いのではないかと思っています。

nanobindとは？

いよいよ本記事の主役であるnanobindを紹介します。
nanobindはpybind11の作者が今年の初め頃から進めている個人プロジェクトで、C++17を使用してpybind11の特徴を引き継ぎつつ、次のような改良が行われています。

パフォーマンスの向上

pybind11よりもバイナリサイズの縮小、コンパイル時間、実行時間の高速化がされてますという話です。
以下のグラフは実行時間に関するグラフなのですが、オレンジがpybind11、緑がnanobindでどれくらい向上しているかが分かります。

テンソルとの親和性の向上

pybind11ではEigenをnumpyの行列に変換してくれるという便利機能があったのですが、nanobindではそれが無くなりました。(最新バージョンではEigenへの変換も入りました！)
代わりにtensorクラスがnanobind内に追加され、それがnumpyやpytorchのテンソルとの変換を行ってくれるようになっています。
これにより、CPUだけでなくGPUやTPUを用いたテンソルとの親和性が向上しました。

CMakeとのインターフェース

pybind11よりも簡単なCMakeとのインターフェースが用意されています。
作者が作ったexampleにCMakeとscikit-buildを使用したプロジェクト例が公開されています。

https://github.com/wjakob/nanobind_example

nanobindを使ってみる

それでは実際にnanobindを使ってみようと思います。
今回のお題は画像上の境界ボックスのIoUを求めるパッケージを作っていきます。Faster-RCNNとかで使うアレですね。

元ネタはCythonを使って実装されたこちらのものでこれをnanobindで作り直すような形になります。
作成したものはnanobboxとしてリポジトリを公開しています。
https://github.com/neka-nat/nanobbox

ビルドまわりの設定

poetryを使って設定する場合、ビルド用にいろいろと追加パッケージが必要になります。
poetryのbuild-systemの項目を以下のように記述します。

[build-system]
requires = [
    "poetry-core>=1.0.0",
    "setuptools>=42",
    "wheel",
    "scikit-build==0.14.0",
    "cmake>=3.18",
    "nanobind>=0.0.3",
    "ninja; platform_system!='Windows'"
]

build-backend = "setuptools.build_meta"

IoU計算

実際のIoU計算部分は以下のような形になります。
tensorクラスを入出力にした関数をC++で実装することで、python側からnumpyで呼ぶことができるようになります。
今回はさらなる高速化のためにOpenMPも使っています。

using bbox_array = nb::tensor<nb::numpy, float, nb::shape<nb::any, 4>, nb::c_contig, nb::device::cpu>;
using cpu_matrix = nb::tensor<nb::numpy, float, nb::shape<nb::any, nb::any>, nb::c_contig, nb::device::cpu>;

cpu_matrix bbox_overlaps(bbox_array& boxes, bbox_array& query_boxes) {
    float *data = new float[boxes.shape(0) * query_boxes.shape(0)] { 0 };
    size_t shape[2] = { boxes.shape(0), query_boxes.shape(0) };
    nb::capsule owner(data, [](void *p) noexcept {
        delete[] (float *) p;
    });
    cpu_matrix overlaps(data, 2, shape, owner);
    #pragma omp parallel for
    for (size_t k = 0; k < query_boxes.shape(0); ++k) {
        const float box_area = (query_boxes(k, 2) - query_boxes(k, 0) + 1) * (query_boxes(k, 3) - query_boxes(k, 1) + 1);
        for (size_t n = 0; n < boxes.shape(0); ++n) {
            const float iw = std::min(boxes(n, 2), query_boxes(k, 2)) - std::max(boxes(n, 0), query_boxes(k, 0)) + 1;
            if (iw <= 0) continue;
            const float ih = std::min(boxes(n, 3), query_boxes(k, 3)) - std::max(boxes(n, 1), query_boxes(k, 1)) + 1;
            if (ih <= 0) continue;
            const float ua = (boxes(n, 2) - boxes(n, 0) + 1) * (boxes(n, 3) - boxes(n, 1) + 1) + box_area - iw * ih;
            overlaps(n, k) = iw * ih / ua;
        }
    }
    return overlaps;
}

バインディング

最後にpythonのバインディングを行う部分です。
pybind11を使ったことのある方は見覚えのある書き方かと思います。

NB_MODULE(nanobbox_ext, m) {
    m.def("bbox_overlaps", &bbox_overlaps,
          nb::raw_doc(
          "Parameters\n"
          "----------\n"
          "boxes: (N, 4) ndarray of float\n"
          "query_boxes: (K, 4) ndarray of float\n"
          "Returns\n"
          "-------\n"
          "overlaps: (N, K) ndarray of overlap between boxes and query_boxes"),
          "boxes"_a, "query_boxes"_a);
}

ベンチマーク

作成したパッケージを使って、実行時間のテストを行いました。
以下のようなコードを用いて計測しています。

import time
import numpy as np
import nanobbox as nb
import cython_bbox as cb

bboxes = np.random.rand(10000, 4)
query_bboxes = np.random.rand(10000, 4)

start = time.time()
result = nb.bbox_overlaps(bboxes, query_bboxes)
print("nanobbox:", time.time() - start)

start = time.time()
result = cb.bbox_overlaps(bboxes, query_bboxes)
print("cython_bbox:", time.time() - start)

結果は以下のようになりました。
OpenMPを使わない状態でも若干ですがパフォーマンスがアップしているのが分かります。

	計算時間
cython_bbox	0.87s
nanobbox	0.77s
nanobbox(omp)	0.22s

まとめ

次世代Pythonバインディングライブラリのnanobindについて紹介しました。
Cythonと比較して高速化ができ、numpyやpytorchなどのテンソルとの親和性も高く、今後pybind11に代わって使われていくのではないかと感じました。
GPUを使った計算のpytorchとの連携なども今後トライしてみたいです。