LibTorch(PyTorch C++API)の使い方

プログラミング

　英語を読める人は素直に公式のドキュメントおよびチュートリアルを読むべき。翻訳ソフトを使ってでもこれらを読んだ方が良いと思う。

　以下の作業は少なくともcuda10.0,cudnn7という環境では動くと思われる。(Dockerでnvidia/cuda:10.0-cudnn7-devel-ubuntu18.04イメージから新しいコンテナを作ってそこで作業した)

簡単なサンプルのビルドまで

　まずダウンロード。以下のページ

から適切にバージョンを指定して落としてくる。

f:id:tokumini:20190824075901p:plain

以前はABIのバージョンがどうこうという問題があって自前でビルドしなければならない場合もあったが(※参考)、今見ると二つ用意されているようなので自分の環境に合わせた方をダウンロードすれば良いと思われる。

　解凍するとlibtorchというディレクトリが得られる。LibTorchを使うときは基本的にcmakeでこのlibtorchディレクトリを指定してコンパイルすることになる。cmakeはCLionでもVisual Studioでも使えるのでそこまで使い勝手は悪くないと思う。

　このページで示されているexample-appというものをコンパイルしてみることとする。example-appというディレクトリを作って、そこにexample-app.cppとCMakeLists.txtという2つのファイルを用意する。

example-app/
    CMakeLists.txt
    example-app.cpp

　example-app.cppは例そのままのtensorを作って表示するというものになる。

#include <torch/torch.h>
#include <iostream>

int main() {
  torch::Tensor tensor = torch::rand({2, 3});
  std::cout << tensor << std::endl;
}

　CMakeLists.txtの方も例とほぼ同じなのだが、一点だけ使いやすいように修正する。というのも、もとの例ではビルドするタイミングでlibtorchのパスを指定しているのだが、いちいちコマンドライン引数にパスを付け加えるのは嫌だし、IDE側から見つけるためにもlibtorchのパスをCMakeLists.txtに書いてしまう。find_package(Torch REQUIRED)の前の行にlist(APPEND CMAKE_PREFIX_PATH ~/libtorch)としてcmakeがパッケージを探すパスを付け加えることができる。今はホーム直下にlibtorchを置いたのでこうしているが、各自ダウンロードしたところへ適当にパスを合わせて書けば良い。

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example-app)

list(APPEND CMAKE_PREFIX_PATH ~/libtorch)
find_package(Torch REQUIRED)

add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 11)

　cmakeではソースコードを置いているディレクトリを汚さないようにbuildというビルド専用のディレクトリをその階層に作るのが普通らしいのでそのようにする。

mkdir build
cd build
cmake ..
make

　これでexample-appを実行すると

 0.8291  0.0013  0.0553
 0.3590  0.1988  0.3392
[ Variable[CPUFloatType]{2,3} ]

と表示された。

ニューラルネットを書く

　ニューラルネットなどの例は公式GitHubののexamplesを見るのが良いと思う。mnistを分類するような簡単なCNNだとmnist.cppにあり、学習の仕方も含めてだいたいはこれで理解できるはず。というわけでこの例で足りない部分についていくらか記述する。といっても正直なところ自分もよくわかっていないところが多いのでむしろ指摘をもらいたい……。

自作モジュールを作る

　ニューラルネットの部分的モジュールなどを書く(たとえばMyModuleというモジュールを作るとする)場合はtorch::nn::Moduleを継承してMyModuleImplというクラスを定義し、最後にTORCH_MODULE(MyModule)というマクロを呼び出すという手順をする(チュートリアルの中程にある)。これによってMyModuleという名前のtorch::nn:ModuleHolderが定義される。Linerについての例だと以下のような感じ。

struct LinearImpl : torch::nn::Module {
  LinearImpl(int64_t in, int64_t out);

  Tensor forward(const Tensor& input);

  Tensor weight, bias;
};

TORCH_MODULE(Linear);

　なぜこんな面倒なことをするかというと、torch::nn::Moduleは大きいクラスになるので、たとえば関数の引数に与えるときに値渡しをすると非効率的になる。C++なのだから参照渡しをすることも容易なんだけど、Python的と同じ書き方で同じ効果を持った方がわかりやすい。Pythonの関数は参照の値渡しなので、C++でそれを実現するにはtorch::nn::Moduleを直接渡すのではなくそれへのstd::shared_ptrであるtorch::nn:ModuleHolderを渡すようにするのが自然。英語があまり読めなかったので適当なことを言っているかもしれない。

モデルをセーブ・ロードする

　torch::nn:ModuleHolderであるnnという変数があったとき、

torch::save(nn, "path_to_model_file");
torch::load(nn, "path_to_model_file");

とする。先のチュートリアル中に

the serialization API (torch::save and torch::load) only supports module holders (or plain shared_ptr).

とあるので、試したことはないがtorch::nn::Moduleはセーブできないのだと思われる。サンプルのmnist.cppとかはtorch::nn:ModuleHolderを使わず書いているのでセーブできなさそうなのだが、そんなものをexampleとして出すのはどうなんだという気もする……。

計算結果を取り出す

　そもそもモデルを定義した後は一度model.to(device);(mnist.cppの129行目)のようにしてGPUにモデルを転送するという作業が必要になる。入力、教師データも同様にauto data = batch.data.to(device), targets = batch.target.to(device);としてGPUに飛ばしてから計算を走らせる。結果を取り出す際には計算結果を示すtorch::Tensorを一度CPUに戻してくる。その後torch::Tensorのメソッドdata<T>()を使って、型Tに対するポインタとして先頭のポインタを得る。自分のコードでは、たとえば1バッチについてDATA_DIM次元のベクトルを計算するデータを取り出す場合は以下のような感じにしている。

    //計算結果を格納するstd::vector
    std::vector<std::vector<float>> results(batch_size);

    //計算してCPUに持ってくる
    torch::Tensor nn_result = model->forward(input).cpu();

    //データの先頭へのポインタを取得してstd::vectorにバッチごとに突っ込む
    float* p = nn_result.data<float>();
    for (uint64_t i = 0; i < batch_size; i++) {
        results[i].assign(p + i * DATA_DIM, p + (i + 1) * DATA_DIM);
    }

　LibTorchどうこうという話ではなくC++としてデータをstd::vectorへと変換するところで非効率的なことをやっている気がする。C++よくわからん。

その他

　一応拙作のMiacisという将棋ソフトでもAlphaZero的なResidual Blockを持つCNNを実装しているため、参考になれば。

　ニューラルネット部分はだいたいsrc/neural_network.cppにある。