PyTorchで学習したモデルをC++から使う

はじめに

この記事では，PyTorch(Python)で学習したモデルをLibTorch(C++)から使う方法について紹介します．
基本的な内容は公式のドキュメントに書いてあります．しかし，学習済みモデルをC++から使うに当たって実践的に必要となる情報の全てがそこに記述されているわけではなかったため，ここに情報をまとめることにしました．
本記事では基本的な内容を超えて，その実践編としてGitHubに公開されているPyTorchを用いた実装を実際にC++で動かすまでの手順をまとめます．
これまで幾つかの学習済みモデルをC++で動かしてきましたが，今回はDepthNetに焦点を当ててPyTorchで学習させたモデルをC++から扱うためのワークフローをまとめます．
実装を読んだほうが早いという方は，こちらへどうぞ．( オレオレ英語でごめん，でも実装は動くから許して )

検証済み環境

OS : Ubuntu 18.04, 16.04
Python : 3.6, 3.7
PyTorch : 1.0.0
OpenCV : 3.4
g++ : 7.3, 6.5, 5.4
GPU/CUDA : 未検証(CPUモードのみ)

PyTorchで学習したモデルをC++で読み込むまでの手順

公式のドキュメントによると以下の手順が必要となります．

1. 学習済みモデルをTorch Script Moduleに変換する
2. Torch Script Moduleを保存する
3. C++から読み込む

では，早速DepthNetにこの手順を当てはめていきましょう．

1. 学習済みモデルをTorch Script Moduleに変換する

学習済みモデルをTorch Script Moduleに変換する方法は，torch.jit.traceを使う方法とAnnotationを使う方法の2通りがありますが，今回はtorch.jit.traceを採用します．
DepthNetの推論時の実装を参考にTorch Module Scriptを保存するスクリプトを書きます．

# 実装例
device = torch.device("cpu")
weights = torch.load(args.pretrained) # args.pretrained -> 学習で得た重みへのパス
depth_net = DepthNet(batch_norm=weights['bn'],
                     depth_activation=weights['activation_function'],
                     clamp=weights['clamp']).to(device)
depth_net.load_state_dict(weights['state_dict'])
# テスト (重み固定)
depth_net.eval()

# 入力するtensorの次元は本家から調べて合わせる
# 次元さえ合っていれば何でも良い．
# input-tensor-size = 1 x 6 x h x w
h = args.img_height #512
w = args.img_width  #512
traced_net = torch.jit.trace(depth_net, torch.rand(1, 6, h, w).to(device))

2. Torch Script Moduleを保存する

1行です．Warningが出ると思いますが，ネットワークに入力するTensorのサイズを，Export時と推論時で揃えている限り問題ないはずです（要検証

traced_net.save("DepthNet_h{}_w{}_{}.pt".format(h,w,mode))
print("DepthNet_h{}_w{}_{}.pt is exported".format(h,w,mode))

3. C++から読み込む

LibTorchを用いたC++プログラムをコンパイルするにあたってLibTorchをダウンロードし，CMakeの設定を行う必要があります．手順はこちら．
読み込みは一行です．

// s_model_name -> torch srcipt module へのパス
std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(s_module_name);

以上で下準備は終わりです．この後はC++を書いていきます．

Torch Script ModuleをC++で実行する手順

最も単純な実装例がこちらです．

// 入力となるat::Tensorを作る [1 x 6 x 512 x 512]
at::Tensor input = torch::ones({1, 6, 512, 512}));
// Moduleを実行し，出力を受け取る
at::Tensor output = module->forward({input}).toTensor();
// 出力Tensorの次元を確認 (DepthNetだと [1 x 128 x 128])
std::cout << "dim = [" << output_tensor.size(0)
          << " x "     << output_tensor.size(1)
          << " x "     << output_tensor.size(2)
          << "]\n";

勿論，入力が torch::onesでは意味がありません．今後は，この入力のat::TensorもDepthNetと同様になるよう作っていきます．
ここで紹介する内容は，参考URLにも記載したPyTorchのissueスレッドを参考としています．

1. 入力となるat::Tensorを作る.

DepthNetはRGB画像を2枚を重ねたTensorを入力としています．
まずOpenCVで画像を読み込み，at::Tensorに変換します．実装例は以下のとおりです．

// 入力画像をペアで読み込む
auto pm_images = std::make_pair(cv::imread(s_image_name0, 1), cv::imread(s_image_name1, 1));

// Moduleに入力するat::Tensorの次元を定義
const int channel = pm_images.first.channels();
const int height = pm_images.first.rows;
const int width = pm_images.first.cols;
std::vector<int64_t> dims{static_cast<int64_t>(1), // 1
                          static_cast<int64_t>(channel * 2), // 6
                          static_cast<int64_t>(height), // h=512
                          static_cast<int64_t>(width)}; // w=512

// 入力Tensorに変換する予定のcv::Matを作っておく．
// cv::Matのサイズも入力のat::Tensorに合わせておく
cv::Mat mf_input = cv::Mat::zeros(height*channel*2, width, CV_32FC1);

// OpenCVではRGBではなくBGRの順に値が入っているので，入れ替える
// DepthNet本家がやっているような下準備もここで行う．
cv::Mat m_bgr[channel], mf_rgb_rgb[channel*2];
cv::split(pm_images.first, m_bgr);
for(int i = 0; i < 3; i++) {
  m_bgr[i].convertTo(mf_rgb_rgb[2-i], CV_32FC1, 1.0/255.0, -0.5); 
  m_bgr[i].release();
}

cv::split(pm_images.second, m_bgr);
for(int i = 0; i < 3; i++) {
  m_bgr[i].convertTo(mf_rgb_rgb[5-i], CV_32FC1, 1.0/255.0, -0.5); 
  m_bgr[i].release();
}

// at::Tensorに変換予定のcv::Matに値をコピーしていく．
for (int i = 0; i < 6; i++) {
  mf_rgb_rgb[i].copyTo(mf_input.rowRange(i*height, (i+1)*height));
}

// 下準備．DepthNetがこうやってる．
mf_input /= 0.2;

// 入力となるat::Tensorを生成
at::TensorOptions options(at::kFloat);
at::Tensor input_tensor = torch::from_blob(mf_input.data, at::IntList(dims), options);

以上で，入力となるat::Tensorを作ることができました．

2. Moduleを実行し，出力を受け取る

at::Tensor output_tensor = module->forward({input_tensor}).toTensor();
// output_tensorの次元 -> [1 x 128 x 128]

3. 出力のat::Tensorからデプスカラーマップを作る

ここまで来たら残りは後始末です．DepthNet本家が行っている処理をC++に書き直して奥行きのカラーマップを作ります．

// 出力のat::Tensorをcv::Matへ格納
cv::Mat m_output(cv::Size(output_tensor.size(1)/*128*/, output_tensor.size(2)/*128*/), CV_32FC1, output_tensor.data<float>());
cv::Mat m_depth = m_output.clone();
cv::Mat m_upscaled_depth;
cv::resize(m_depth, m_upscaled_depth, cv::Size(width/*512*/, height/*512*/), 0, 0);
const float max_value = 100.0; // DepthNet本家が決めてた
cv::Mat m_arranged_depth = 255.0*m_upscaled_depth/max_value; 
for (int i = 0; i < m_arranged_depth.rows; i++) {
  for (int j = 0; j < m_arranged_depth.cols; j++) {
    if(m_arranged_depth.at<float>(i,j) > 255.0) {
      m_arranged_depth.at<float>(i,j) = 255.0;
    } 
    else if (m_arranged_depth.at<float>(i,j) < 0.0) {
      m_arranged_depth.at<float>(i,j) = 0.0;
    }
  }
}
m_arranged_depth.convertTo(m_arranged_depth, CV_8U); 
cv::Mat m_color_map;
cv::applyColorMap(m_arranged_depth, m_color_map, cv::COLORMAP_RAINBOW);

結果

元々の入力画像と奥行きのカラーマップです．本家と同様な結果が得られました．

終わりに

PyTorchのC++APIは未だ発展途上です．そのため，幾つか注意しておくべき点があります．
ひとつ目は，今後C++APIに何らかの破壊的変更が加わってここで紹介した内容が参考にならなくなる可能性です．
ふたつ目は，同じ学習済みモデルであっても，C++から使った場合とPythohから使った場合とで得られる結果が完全には一致しないという点です．厳密な評価に手を付けていないのですが，Pythonから実行した場合と完全に同じ出力を期待すると痛い目を見るかもしれません．
しかし，上記の2点を目を瞑ればPyTorchのC++APIは十分実用的です．一度使い方を把握すれば大凡のモデルに対して適応することができるため，C++の実装に深層学習を用いた何かしらを組み込みたいとなったときに非常に有用であることには変わりません．
今回はGPU/CUDAを使った計算については検証していないので，引き続き調べてみようと思います．
慣れない技術記事を長々と書いてしまいましたが，結局のところ試してもらうのが一番早いので，こちらをどうぞ．使い方はREADMEに書いてあります．
https://github.com/cashiwamochi/DepthNet

参考URL

DepthNet (by Clément Pinard)
https://github.com/ClementPinard/DepthNet
PYTORCH C++ API
https://pytorch.org/cppdocs/
LOADING A PYTORCH MODEL IN C++
https://pytorch.org/tutorials/advanced/cpp_export.html
INSTALLING C++ DISTRIBUTIONS OF PYTORCH
https://pytorch.org/cppdocs/installing.html
PyTorch issue#12506
https://github.com/pytorch/pytorch/issues/12506