この記事はkaggle Advent Calendar 2021の2枚目の20日目の記事です。
本記事では、Riiidコンペ（もう一年前ですが…）でC++のpybind11で拡張モジュールを書いてメモリ使用量を節約した話をします。（自分の最終順位は51位でした）pybind11で拡張モジュールをどう書くかの入門にもなっていますので、何となくCやC++をPythonでどう使うのか興味がある方も読んでいただければと思います。
構成としては二部構成で、まず前半（実践編）でRiiidコンペがどんなものでどのように拡張モジュールを使ったかをご紹介して、後半（理論編）でCPythonが拡張モジュールをどのようにimportしているかの話をします。
注意1
以下、kaggle notebookに関係のあるLinux OSに話を限定します。また、Pythonのインタプリタにはいくつかの言語で実装が存在して、その中でもC言語で書かれたCPythonは広く使われているだけでなく、Pythonのリファレンス実装としての役割も持っていると言われており、実際Python公式ドキュメントの色々なところでCPythonの実装の詳細についての言及があります。普段kaggle notebookで使っているPythonの実装もCPythonです（もちろん使おうと思えば他実装も使えますが）ので、以下はCPythonに話を限ります。
注意2
理論編ではCPythonでgdbを使って拡張モジュールが実際にimportされるところを確認します。CPythonのコードを見るとわかりますが、環境によって実装が違うので、デバッガの出力や挙動が異なる場合があります。例えばdynamic loadingはWindowsでは本文で紹介するPython/dynload_shlib.c内のdlopenではなくPython/dynload_win.c内のLoadLibraryExW関数で行われます。自分の実行環境はWSL2 Ubuntu 20.04.3 LTSで、この場合はdlopen関数が使われます。2021/12/19現在、kaggle notebookの実行環境は同じくUbuntu 20.04.03 LTSです。
実践編：Riiidコンペで拡張モジュールをどう使ったか
Riiidコンペの概要
Riiidコンペはオンラインの学習プラットフォームで集められたデータを使ったコンペで、学生の過去の講義視聴や問題解答のデータを学習データとして、次に見る問題に正解したかどうかを予測するコンペでした。特徴的な点として、テストデータの予測が、notebookを提出してそのnotebook上で
- 少量（数個～数千個程度）のデータについての予測
- 予測したデータの正解が与えられたうえで、時系列の次の少量データについて予測
- 以降は2の繰り返し
といった形の、KaggleのTime-series APIを使ったオンライン予測をするものになっていたことが挙げられます。
予測の手段としては、ユーザIDと問題IDがデータのラベルとして与えられるので、データをユーザ毎の時系列データに持ち直して、ユーザが過去に見た問題の正解・不正解データから今見ている問題の正解・不正解を予測するTransformerを使う解法が上位解法には多かった印象です。（自分もこの手法を取りました）
もう一つの特徴として、学習データ、予測データともにデータ量が大きいことが挙げられます。学習データのcsvは約1億行あり、予測データの数は約250万個あります。このため、時間とメモリの制限が問題になります。
時間制限の問題
コンペのルールによるとGPU付きのnotebookで予測が9時間以内に終わらなければならないのですが、バッチ予測ではなくオンライン予測で、しかもTransformerを使った予測をしようとするとかなりカツカツになります。
この点上位解法には面白い工夫があって、自分が把握している範囲（工夫があるのにここに挙げていなかったら自分が把握していないというだけです🙇）だと、3位のチームの解法では、テストデータの最初の20%がpublicで残り80%がprivateなので、はじめの20%については全く予測を走らせないで時間を節約することで、残りの80%で時間をフルに使って2つのTransformerの予測のアンサンブルを取ったようです。（もちろんpublic scoreがめちゃくちゃに低くなるのでprivate scoreの開示までどうなったかわからないということで、このアプローチにThe Blindfolded Gunslingerという名前をつけているようです）
また、1位のSeungKeeの解法は(コンテストの時はkeetarという名前でした)、TransformerのAttentionのQK^T積のQ（クエリ行列）に入れるデータとして、予測する部分である時系列データの最後のデータのみを入れてもスコアがあまり変わらないことを見つけて、予測時のAttention部分の演算をインプットの長さに対して二乗から線形に落とすことで、Transformerのインプットデータの長さを1728まで伸ばしてかつ5つのモデルでアンサンブルするというアプローチをとっています。（自分はインプットのデータの長さを90にしてもかなり時間を使ったので1728というのは相当な長さです、しかもモデル5つ…）
メモリ制限の問題
本記事の本題はこちらです。オンライン予測なので、ユーザ特徴量を予測時に与えられる正解データに応じて更新しながらdictで持っておきたくなりますが、例えば各ユーザが各問題を何回見たことがあるかという特徴量を学習データから作ると、キーの数が86867031個の巨大なdictが生成されます。GPU付きのkaggle notebookのメモリは13GBしかありませんので、Pythonのdictで持つと、いくつかユーザ特徴量を使おうとするとメモリが足りなくなります。
というわけでメモリ効率のよいdict、以降hashmapと呼びます、を使いたくなります。Pythonのライブラリから探してもいいんですが、実はpybind11というライブラリを使うとC++で書かれたプログラムとPythonとを手軽に連携できるようになりますので、C++のライブラリから探します。C++の標準ライブラリ(STL)に入っているunordered_mapと呼ばれるhashmapはメモリ効率があまり良くないので、ここではparallel_hashmapというメモリ効率の良いhashmapを使います。
以下、どのようにpybind11を使うかコードで説明していきます。kaggle notebook上で動く例はこちらのnotebookにありますので、実際に使ってみたい方はそちらをご参照ください。(GPUを使うのでQuotaにご注意ください）
pybind11でparallel_hashmapを使う
まずC++でparallel_hashmapについての必要なメソッドをまとめたstructを定義します（classでもよいです）。parallel_hashmapはヘッダーをincludeするだけで使えるライブラリで、まずは次のようにライブラリをincludeしてstructのメンバー変数で使います（parallel_hashmapにはいくつかのhashmapのデータ構造が入っているのですが、ここではflat_hash_mapを使います）：
#include <parallel_hashmap/phmap.h>
using phmap::flat_hash_map;
struct riiid_user_x_content_int {
flat_hash_map<uint32_t, flat_hash_map<uint16_t, uint8_t>> m_dict;
};Riiidコンペでは（ユーザー×問題）毎の特徴量を扱いたいので、ネストしたhashmapを使って、m_dict[user_id][content_id]（content_idは問題ID）に特徴量を入れます。Riiidコンペのuser_idとcontent_idの範囲から、適切に32,16,8bitの非負整数型uint32_t,uint16_t,uint8_tを使ってメモリを節約します。特徴量として実数のものを使いたい場合はuint8_tの部分をdoubleなどに適宜変更します。
ここで、pybind11に渡しやすいようにgetやsetのメンバー関数を定義します：
#include <parallel_hashmap/phmap.h>
using phmap::flat_hash_map;
struct riiid_user_x_content_int {
int getval(int user_id, int content_id) { return m_dict[user_id][content_id]; } void setval(int user_id, int content_id, int value) { m_dict[user_id][content_id] = value; }
flat_hash_map<uint32_t, flat_hash_map<uint16_t, uint8_t>> m_dict;
};次に、実際kaggleコンペで使うときには学習段階で特徴量を作っておいて提出したコード上で予測時にロードしたいので、このstructの初期化関数（コンストラクタ）に特徴量データファイルへのファイルパスを渡せばm_dictに特徴量を入れてくれるようにしましょう。ここではcsvファイルで保存したデータをロードするようにします。（少々長く本質的な部分ではないので必要に応じて読み飛ばして下さい。（今回は自前で書いてしまいますが、一般にC++でcsvをパースするときはエッジケースへの対応などが必要なので既存のライブラリを使うことをおすすめします。））
#include <parallel_hashmap/phmap.h>
#include <string>#include <fstream>#include <iostream>
using phmap::flat_hash_map;
struct riiid_user_x_content_int {
riiid_user_x_content_int(const std::string& datapath) { std::ifstream mycsvfile (datapath); if (mycsvfile.is_open()) { std::string line; getline(mycsvfile, line); while (mycsvfile.good()) { getline(mycsvfile, line); if (line.size() == 0) break; int comma_pos = line.find(',', 0); int comma_pos2 = line.find(',', comma_pos+1); uint32_t user_id = stoi(line.substr(0, comma_pos)); uint16_t content_id = stoi(line.substr(comma_pos+1, comma_pos2-comma_pos-1)); uint8_t value = stoi(line.substr(comma_pos2+1)); m_dict[user_id][content_id] = value; } std::cout << "csv file loaded" << std::endl; } else { std::cout << "Error in opening file" << std::endl; } }
int getval(int user_id, int content_id) {
return m_dict[user_id][content_id];
}
void setval(int user_id, int content_id, int value) {
m_dict[user_id][content_id] = value;
}
flat_hash_map<uint32_t, flat_hash_map<uint16_t, uint8_t>> m_dict;
};これで、parallel hashmapを使った特徴量を収納するクラスができました。あとはPythonから扱えるようにするだけです。これはpybind11を使うと非常に簡単にできます：
#include <pybind11/pybind11.h>
namespace py = pybind11;
PYBIND11_MODULE(riiid_module, m) {
py::class_<riiid_user_x_content_int>(m, "user_x_content_int")
.def(py::init<const std::string &>())
.def("get", &riiid_user_x_content_int::getval,
py::arg("user_id"), py::arg("content_id"))
.def("set", &riiid_user_x_content_int::setval,
py::arg("user_id"), py::arg("content_id"), py::arg("value"));
}PYBIND11_MODULEマクロでは、PythonからどのようにC++の関数やクラスを呼び出すかを定義します。riiid_moduleがモジュールの名前になります。ここではモジュール内に、Pythonから扱うクラス名はuser_x_content_intとして、実装はC++で書いたriiid_user_x_content_intとしたクラスをpy::class_を使って定義します。その他の初期化関数、get、setは関数ポインタやpy::arg関数を使ってpy::class_上でdef関数を使って定義していきます。（詳細はドキュメントをご参照ください）
このように定義したら、上の二つのコードをまとめてhoge.cppなどとして保存して、共有ライブラリ（.soファイル）にコンパイルします（コンパイルオプションはpybind11のドキュメントに従います）：
export CPP_FILEPATH="/path/to/hoge.cpp"
export HASHMAP_PATH="/path/to/hashmap_folder"
c++ -O3 -Wall -shared -std=c++11 -fPIC `python3 -m pybind11 --includes` $CPP_FILEPATH -I$HASHMAP_PATH -o riiid_module`python3-config --extension-suffix`これによって{モジュール名}.{CPythonのバージョンやマシンの情報}.soファイルができます。（例：riiid_module.cpython-38-x86_64-linux-gnu.so）この状態で、以下のようにPythonからriiid_moduleモジュールをimportすれば.soファイルを読んでくれて、クラスが使えるようになります。
from riiid_module import user_x_content_int
path_to_csv = "path/to/csv"
cumulative_count = user_x_content_int(path_to_csv)
cumulative_count.set(user_id=42, content_id=4, value=3)
assert cumulative_count.get(user_id=42, content_id=4) == 3 # キーワード引数
assert cumulative_count.get(42, 4) == 3 # ポジション引数parallel hashmapでRiiidコンペの(ユーザ×問題)毎の特徴量を持ったときのメモリ使用量のベンチマークはこちらで、特徴量が16bit整数のとき700MiB、32bit実数のとき1300MiB程度でした。Pythonのfloat64のdictで8.6GiB程度なので、これと比べてメモリの節約になっていることがわかります。コンペ中に困った現象としてpickleをloadするときにメモリ使用量が制限を超えてしまってnotebookが落ちるというのがあったのですが、C++でcsvを少しずつbufferでロードしていけばこのような心配もありません。
このように、pybind11を使ってPythonからクラスをどう扱うかDSLで定義してやれば、あとはPythonとC++をつなぐC APIの部分はpybind11がテンプレートプログラミングで何とかしてくれるので、意外とC++で拡張モジュールを書くハードルは低いです。この裏でどのようなことが起きているかを知りたい方は続けて理論編に進んでいただいて、あまり気にならない方はまとめと展望まで飛んでいただければと思います。
理論編：拡張モジュール、共有ライブラリ、dlopen関数
実践編でC++でコードを書いて、pybind11を使ってPythonからモジュールとして呼び出せるようにする流れを見てきました。このようなC/C++で書かれたモジュールは拡張モジュールと呼ばれます。拡張モジュールが何のためにあるかPythonの公式ドキュメントから引用すると、
It is quite easy to add new built-in modules to Python, if you know how to program in C. Such extension modules can do two things that can’t be done directly in Python: they can implement new built-in object types, and they can call C library functions and system calls.
というわけで、Pythonから呼び出すモジュールの実装の中身はCでも書けるので、パフォーマンスが重要な型をCで実装したいとき、Cのライブラリ関数を使いたいとき、またシステムコールでOSとやり取りしたくなったりなったときに使う、といった感じです。他にも、既存のライブラリでCから呼び出せるようになっている（CのAPIが用意されている）ものはたくさんあるので、それらのライブラリをCから呼び出すAPIとPythonからCを呼び出すAPIを適応させる薄いラッパーを書けばライブラリが使いまわせる、といった重要な利点もあります。
拡張モジュール自体の目的としてははっきりした理由があることがわかったので、次に拡張モジュールの仕組みについて見ていきます。
以下、dynamic loadingなどといった話題が出てきますが、.soファイルなどのライブラリファイルを静的または動的にリンカでリンクする話については、lldやmoldといったリンカを開発されている専門家中の専門家の植山類さんの講演「リンカの基礎知識」がYouTubeにあるのでそちらをおすすめします。
インタプリタを拡張するとはどういうことか
Pythonの公式ドキュメントに、Extending and Embedding the Python Interpreterというタイトルのページがあり、以下のような記述があります:
This document describes how to write modules in C or C++ to extend the Python interpreter with new modules.
ここのmoduleというのは拡張モジュールのことです。インタプリタをモジュールで「拡張する」というのはどういうことでしょうか。まず、インタプリタにはC言語で書かれた実体のプログラムがあって、CPythonという名前が付いていたことを思い出しましょう。CPythonのコンパイル時に一緒にモジュールを追加してしまってもよいのですが(ModulesフォルダにCで書いた実装を追加してLibフォルダにPythonとのinterfaceを追加する感じでしょうか)、それだとモジュールを追加するたびにCPythonを再コンパイルする必要がありとても大変なので、ここでは既にコンパイルされたCPythonを動的に拡張することを考えます。
一般にプログラムを動的に拡張する方法の一つとして、dynamic loadingがあります。dynamic loadingとは、共有ライブラリと呼ばれるファイル(拡張子は.soで、実践編でコンパイルして出来たファイルがこれです)の中身を動的リンカ/ローダがメモリ上に配置して、それらのアドレスをメインのプログラムに教えて使えるようにすることを指します。
Cではdynamic loadingのAPIとしてdlopenと呼ばれるライブラリ関数が用意されています。dlopenに.soファイルへのパスを渡すと、共有ライブラリがメモリ上にロードされ、共有ライブラリの関数などのロードされたメモリ上の位置のアドレス情報を含んだhandleオブジェクトが返され、これを通じて関数などをメインのプログラム(ここではCPythonのこと)から使えるようになります。このようにして、メインのプログラムとは別にコンパイルされたプログラムが動的に利用できます。
CPythonにおいても、(使える環境では)このdlopen関数によってPythonインタプリタの拡張ができるようになっています。（コードでいうとこの部分です）
というわけで共有ライブラリを使えばよいことが分かったわけですが、CやC++で書いたプログラムの型や関数をPythonから扱えるようにしたいので、Pythonの言語仕様のPython/C APIを使って拡張モジュールを書く必要があります。ここで、頑張ってPython/C APIに従ってC/C++のコードを書いてもよいのですが、pybind11を使うとマクロやヘルパー関数を使って分かりやすい形で定義できるので、実践編で見たように簡単に書くことができます。実際のところ直接C APIを使うのではなく、このようなthird party toolを使うことが公式のドキュメントでも推奨されています：
This section of the guide covers creating C and C++ extensions without assistance from third party tools. It is intended primarily for creators of those tools, rather than being a recommended way to create your own C extensions.
ここから、実際にCPythonがdlopen関数でdynamic loadingを行うのをデバッガ（gdb）で確認していきます。少々準備が必要なので、確認しなくても十分であればまとめと展望まで飛んでください。
拡張モジュールのimportにdlopenが使われていることの確認
では、CPythonのdlopenを開いているところでgdbのbreakpointを設定してみましょう。以下はWSL2のUbuntu 20.04での実行例で、CPythonのバージョンは3.9とします。
まずCPythonをダウンロードしてきます。自分の環境と揃えたい場合は以下のように3.9ブランチの特定コミットにcheckoutしてください。
git clone --branch 3.9 https://github.com/python/cpython
cd cpython
git checkout 6b0ea06e続けて、ビルド手順に従って以下のようにCPythonをビルドします。まず、依存パッケージのインストールを行います。ここでは以下のようにインストールします。
sudo apt install build-essential
sudo apt install libssl-dev zlib1g-dev libncurses5-dev \
libncursesw5-dev libreadline-dev libsqlite3-dev libgdbm-dev \
libdb5.3-dev libbz2-dev libexpat1-dev liblzma-dev libffi-dev次にCPythonをビルドします：
./configure --with-pydebug
make -j4 -sコマンドが完了したらPythonインタプリタのexecutableファイルpythonが生成されて、以下のように実行できるようになります。
./python次にpybind11をCPythonのsubmoduleとしてpybind11を追加して、PYTHONPATHを通します。
git submodule add -b stable ../../pybind/pybind11 extern/pybind11
git submodule update --init
export PYTHONPATH="$PYTHONPATH:$PWD/extern/pybind11"ここでテスト用の簡単な拡張モジュールを書いてみます。以下のようにpybind11でexampleというモジュールにaddという関数を定義します：
#include <pybind11/pybind11.h>
int add(int i, int j) {
return i + j;
}
PYBIND11_MODULE(example, m) {
m.def("add", &add, "A function which adds two numbers");
}これをexample.cppというファイル名で保存して以下のように.soファイルにコンパイルします。
c++ -O3 -Wall -shared -std=c++11 -fPIC `./python -m pybind11 --includes` example.cpp -o example.cpython-39-x86_64-linux-gnu.soここでpybind11モジュールが見つからないというエラーが出たらPYTHONPATHが通っているかチェックしてください。また、.soファイルについて、自分の環境ではx86_64系、linux gnu用のsuffix（.cpython-39-x86_64-linux-gnu.so）ですが、ここはマシンに応じて適宜変えます。（自分はpython3-config --extension-suffixコマンドの出力のCPythonバージョンだけ変えました）
実行が完了したらexample.cpython-39-x86_64-linux-gnu.soファイルができて、以下のようにCPythonから使えるようになります：
./python -c "from example import add; print(add(2,3))"このようにすると、先ほどビルドした.soファイルのexampleモジュールからadd関数をインポートして、5を標準出力に出してくれます。ようやく拡張モジュールをロードできるようになったので、次にデバッガを使ってCPythonがどのようにdynamic loadingを行っているか調べてみましょう。ここではgdbデバッガを使います。まず、以下のように~/.gdbinitに設定を追加してgdbの自動ロードをオンにします：
add-auto-load-safe-path /path/to/cpython次に、gdbを実行して、dlopen関数を使っている箇所にbreakpointを設定して、import exampleを実行してみます：
gdb ./python
(gdb) b Python/dynload_shlib.c:101
(gdb) run -c "import example"すると、breakpointで実行が停止してその時点での変数やcall stackが見られるようになります：
Starting program: /home/takkyu/cpython/python -c "import example"
[Thread debugging using libthread_db enabled]
Using host libthread_db library "/lib/x86_64-linux-gnu/libthread_db.so.1".
Breakpoint 1, _PyImport_FindSharedFuncptr (
prefix=0x555555877605 "PyInit",
shortname=shortname@entry=0x7ffff7755a20 "example",
pathname=pathname@entry=0x7ffff7766ae0 "/home/takkyu/cpython/example.cpython-39-x86_64-linux-gnu.so",
fp=fp@entry=0x0) at ./Python/dynload_shlib.c:102
102 if (handle == NULL) {ここで表示されているpathname変数がdlopenで開いているファイルパスです。コンパイルした共有ライブラリを読みに行ってくれていることがわかります。btコマンドでcall stackを表示してみるとどのようにここまでたどり着いたのか分かって面白いのですが、ここでは割愛します。
一方で、CPythonと一緒にコンパイルされている標準モジュールをimportしてみます。
gdb ./python
(gdb) b Python/dynload_shlib.c:101
(gdb) run -c "import sys, os"これらはdynamic loadされていないのでbreakpointには引っかかりません：
Starting program: /home/takkyu/cpython/python -c "import sys, os"
[Thread debugging using libthread_db enabled]
Using host libthread_db library "/lib/x86_64-linux-gnu/libthread_db.so.1".
[Inferior 1 (process 22099) exited normally]というわけでPythonが拡張モジュールのimportにdlopen関数を使っていることが確認できました。
まとめと展望
本記事ではpybind11で拡張モジュールを使うことについてご紹介しました。PythonにC
APIがあることの理由の一つとして既存のライブラリのラッパーを書くだけでPythonから手軽に使えるようになる、というものがありますが、自分がRiiidで使ったコードもparallel_hashmapの[]演算をgetとsetに書き直しただけなのでこの用途に当たります。イメージとしてはPythonがフロントエンド、バックエンドはCで書かれたライブラリで、Pythonが複数のバックエンド（parallel_hashmapとTensorflow）の間の通信のハブとしても使えるといった感じでしょうか。
今回はparallel_hashmapから取り出すデータは単一の整数や実数で、ここがボトルネックにならなかったのでpybind11が（implicitに）提供する型変換で事足りましたが、より大量の連続するデータを扱いたい場合はbuffer protocolやmemory viewといったトピックが関わってきます。（pybind11の該当ドキュメントはこのあたりです）また、Pythonのガベージコレクション（参照カウント）とC++のメモリ管理を共存させるためにはpybind11の提供するreturn value policiesといった機能が必要になります。
このような難しさがあるということは念頭に置きつつ、pybind11のおかげで拡張モジュールを書くのはそこまで大変ではないということが伝わって、困ったときの一つの選択肢として頭の片隅に置いておいて頂けるようになれば本記事の目的は達成されたと言えます。