PyTorch の C++ API LibTorch 入門 - ネコと和解せよ

LibTorch is 何

PyTorch の C++ APIで, C++でPyTorchライクに機械学習を行うためのライブラリです。他にもPythonで学習したモデルをC++で呼び出して使うといったことができます.

Version1.0の時はサンプルコードが動かなかったり、動いたと思ったら、少しバージョンが上がったらまた動かなくなったりと、割と不安定な状態でしたが、Version1.1である程度安定した（と思うので）備忘録としてまとめておきたいと思います.

LibTorch is 何
Install
とりあえず動かしてみる.
- トラブルシューティング
MNIST サンプルコードを動かす
MNIST サンプルコードの中身
最後に
c.f.

Install

以下から自分の環境にあったものをダウンロードします.
https://pytorch.org/

基本的にこの記事ではMacOSないしはLinux環境を想定してます. ちなみにMacOS版のバイナリファイルはCPU版しかないようです.

ファイルをダウンロードしたら、適当に/usr/local/includeあたりにlibtorchを配置しておきます.

とりあえず動かしてみる.

インストールが完了したので、試しに以下の2行3列の行列を表示させるプログラムを作成する.

$$ \begin{bmatrix} 1 & 1 & 1 \\ 1 & 1 & 1 \end{bmatrix} $$

# main.cpp

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

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

LibTorchを使用する場合にはビルドシステムとしてCMakeを使うことが推奨されている. なのでCMakeLists.txtファイルを作成する.

# CMakeLists.txt

cmake_minimum_required(VERSION 3.0)
project(torchSample)

find_package(Torch REQUIRED)

add_executable(torchSample main.cpp)

target_link_libraries(torchSample "${TORCH_LIBRARIES}")

set_property(TARGET torchSample PROPERTY CXX_STANDARD 11)

# ビルドする

$ mkdir build && cd build

# CMAKE_PREFIX_PATHにLibTorchまでの絶対パスを設定する.
# /usr/local/include/libtorchであれば、
$ cmake -DCMAKE_PREFIX_PATH=/usr/local/include/libtorch ..

$ make

# 実行する

$ ./torchSample

# 出力
# 1  1  1
# 1  1  1
# [ Variable[CPUType]{2,3} ]

とりあえずこれでざっくり動作確認ができたかと思います.

トラブルシューティング

こんな感じのエラーがでたら
Library not loaded: /usr/local/opt/libomp/lib/libomp.dylib
それはOpenMPがインストールされていないということなので、インストールしましょう.

# MacOS
$ brew install libomp

MNIST サンプルコードを動かす

このリポジトリに MNISTのサンプルコードがあります. これをgit cloneして利用します. examples/cpp/mnist/配下にMNIST用のサンプルコードがあります.

ビルドしてみましょう

$ pwd
# => /path/to/example/cpp/mnist

$ ls
# CMakeLists.txt
# README.md
# mnist.cpp

$ mkdir build && cd build

# /path/to/libtorch には libtorchまでの絶対パスを入れます
# 実行中にPythonスクリプトが実行されてMNISTデータセットがダウンロードされます
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..

# ビルド
make

ビルドが完了したら下記コマンドで実行します.

./mnist

実行すると以下のような出力がされるはずです.

Training on CPU.
Train Epoch: 1 [59584/60000] Loss: 0.2291
Test set: Average loss: 0.2118 | Accuracy: 0.934
Train Epoch: 2 [59584/60000] Loss: 0.1293
Test set: Average loss: 0.1304 | Accuracy: 0.959
Train Epoch: 3 [59584/60000] Loss: 0.1772
Test set: Average loss: 0.1071 | Accuracy: 0.965
Train Epoch: 4 [59584/60000] Loss: 0.1150
Test set: Average loss: 0.0894 | Accuracy: 0.972
Train Epoch: 5 [59584/60000] Loss: 0.0998
Test set: Average loss: 0.0824 | Accuracy: 0.974
Train Epoch: 6 [59584/60000] Loss: 0.0497
Test set: Average loss: 0.0706 | Accuracy: 0.977
Train Epoch: 7 [59584/60000] Loss: 0.0502
Test set: Average loss: 0.0688 | Accuracy: 0.979
Train Epoch: 8 [59584/60000] Loss: 0.0674
Test set: Average loss: 0.0624 | Accuracy: 0.981
Train Epoch: 9 [59584/60000] Loss: 0.0529
Test set: Average loss: 0.0598 | Accuracy: 0.981
Train Epoch: 10 [59584/60000] Loss: 0.0545
Test set: Average loss: 0.0554 | Accuracy: 0.984

MNIST サンプルコードの中身

このMNISTサンプルコードについて一部を取り上げて中身を見てみます.

モデル定義

PyTorchでモデル定義するときはtorch.nn.Moduleを継承していたように、 LibTorchでもtorch::nn:Moduleを継承して定義します.

struct Net : torch::nn::Module {
  Net()
      : conv1(torch::nn::Conv2dOptions(1, 10, /*kernel_size=*/5)),
        conv2(torch::nn::Conv2dOptions(10, 20, /*kernel_size=*/5)),
        fc1(320, 50),
        fc2(50, 10) {
    register_module("conv1", conv1);
    register_module("conv2", conv2);
    register_module("conv2_drop", conv2_drop);
    register_module("fc1", fc1);
    register_module("fc2", fc2);
  }

  torch::Tensor forward(torch::Tensor x) {
    x = torch::relu(torch::max_pool2d(conv1->forward(x), 2));
    x = torch::relu(
        torch::max_pool2d(conv2_drop->forward(conv2->forward(x)), 2));
    x = x.view({-1, 320});
    x = torch::relu(fc1->forward(x));
    x = torch::dropout(x, /*p=*/0.5, /*training=*/is_training());
    x = fc2->forward(x);
    return torch::log_softmax(x, /*dim=*/1);
  }

  torch::nn::Conv2d conv1;
  torch::nn::Conv2d conv2;
  torch::nn::FeatureDropout conv2_drop;
  torch::nn::Linear fc1;
  torch::nn::Linear fc2;
};

PyTorchのMNISTサンプルと比較しても似ていることがよく分かるかと思います. 一方でLibTorchではregister_moduleを利用して各モジュール(conv1, conv2, ...etc)を登録するようです.

畳み込み層の初期化には引数に
torch::nn::Conv2dOptions(InputChannel, OutputChannel, KernelSize) を与えています. 特に指定しなければ、stride=1, padding=0となるようです.

データローダー

以下でダウンロードしたデータからMNISTデータセットを読み込み、ノーマライズし、データとラベルでそれぞれ一つのテンソルとなるようにConcatenatesしてます.

auto train_dataset = torch::data::datasets::MNIST(kDataRoot)
                         .map(torch::data::transforms::Normalize<>(0.1307, 0.3081))
                         .map(torch::data::transforms::Stack<>());

上記で読み込んだデータセットをデータローダーに、<>br サンプラーtorch::data::samplers::SequentialSamplerとバッチサイズを指定してセットしています.

auto train_loader =
    torch::data::make_data_loader<torch::data::samplers::SequentialSampler>(
        std::move(train_dataset), kTrainBatchSize);

学習コード

template <typename DataLoader>
void train(
    int32_t epoch,
    Net& model,
    torch::Device device,
    DataLoader& data_loader,
    torch::optim::Optimizer& optimizer,
    size_t dataset_size) {
  model.train();
  size_t batch_idx = 0;
  for (auto& batch : data_loader) {
    auto data = batch.data.to(device), targets = batch.target.to(device);
    optimizer.zero_grad();
    auto output = model.forward(data);
    auto loss = torch::nll_loss(output, targets);
    AT_ASSERT(!std::isnan(loss.template item<float>()));
    loss.backward();
    optimizer.step();

    if (batch_idx++ % kLogInterval == 0) {
      std::printf(
          "\rTrain Epoch: %ld [%5ld/%5ld] Loss: %.4f",
          epoch,
          batch_idx * batch.data.size(0),
          dataset_size,
          loss.template item<float>());
    }
  }
}

同様にPyTorchのMNISTサンプルと比べてもわかるように、非常によく似ています.

The C++ Frontend ではloss.template item<float>()ではなくloss.item<float>()と書いていますが、私の環境では動きませんでした.

最後に

モデルを保存するためには異なる書き方が必要なので、気が向いたら続きを書きたいと思います.

c.f.

PyTorch C++ API — PyTorch master documentation