はじめに
今回はPytorchで作成・学習したモデルをC++で使用する方法を紹介します。
機械学習モデルをC++で利用する用途は次の挙げられます。
- C++で実装された既存のサービスにAIを組み込む
- ROS(Robot Operating System)などPython3に対応していないシステムにPython3用のモデルを組み込む
- そのほかのC++のみが動作するサービスに学習済みモデルを組み込む
それでは実装方法です。
Pytorchでの変換処理
はじめにPythonでの処理について説明を行います。
なお、Pytorchの環境構築は完了しているものとしていますので、ご注意ください。
また、C++で使用したいネットワークもimportできるようにしてください。
ソースコード
次のサンプルコードを実行することでC++で使用するファイルを出力することができます。
# Pytorchのインポート
import torch
# 学習に使用するネットワーク
import Net # 適宜インポートするネットワークを変更してください。
def main():
# 学習に使用したモデルのロード
net = Net()
# 学習済みの重みをネットワークに適用
net.load_state_dict(torch.load("model to path"))
# 学習時の入力サンプル
example_input = torch.rand(1,3,256,256).to("cuda")
# 学習済みモデルのトレース
traced_net = torch.jit.trace(net,example_input)
# 変換モデルの出力
traced_net.save("save to path/traced_model.pt")
if __name__=="__main__"
main()解説
8行目の net = Net() では学習に使用したネットワークをインスタンス化しています。この部分はそれぞれで使用したいネットワークに変更してください。
10行目の net.load_state_dict(torch.load(“model to path”)) では8行目でインスタンス化したモデルに重みを適用しています。“model to path” は学習済みモデルの重みが存在するパスを指定するようにしてください。
12行目の example_input = torch.rand(1,3,256,256).to(“cuda”) では学習したネットワークの入力を仮に作成しています。なお、入力データのチャンネル数と学習時に使用したデータのチャンネル数は一致している必要があります。また、学習時に入力を .to(“cpu”) としている場合はこの入力例も .to(“cpu”) にする必要があるため、ご注意ください。
14行目の traced_net = torch.jit.trace(net,example_input) では入力されたネットワークをトレースします。この処理でネットワークに重みを紐づけてC++で使用できるように変換を行なっています。torch.jit.trace について詳しく知りたい方はこちらをご覧ください。
16行目の traced_net.save(“save to path/traced_model.pt”) では14行目にて変換したモデルを “save to path” に指定したパスに保存します。
実行
先ほどのサンプルコードを記述したPythonファイルをコマンドラインから実行してください。
コンソール上には「 98.1245678 vs 97.12345678 」といった表示がされますが、 traced_model.pt というファイルが指定したパスに出力されていれば問題ありません。
C++編
続いてC++側での処理についてです。
先ほどのPython側の処理で作成した traced_model.pt というファイルを使用していきますので準備をしてください。
また、C++側ではLibtorchというライブラリが必要になりますので、公式サイトよりライブラリをダウンロードしてプロジェクトに入れておいてください。
ソースコード
// Libtorchのインクルード
#include <torch/script.h>
int main()
{
// torch::jit::script::Module 型で module 変数の定義
torch::jit::script::Module module;
// 変換した学習済みモデルの読み込み
module = torch::jit::load("model to path/traced_model.pt");
// モデルへのサンプル入力テンソル
torch::Tensor input = torch::ones({1, 3, 256, 256}).to("cuda");
// 推論と同時に出力結果を変数に格納
auto elements = module.forward(input).toTuple() -> elements();
// 出力結果
auto output = elements[0].toTensor();
return 0;
}解説
2行目では Libtorch のヘッダファイルをインクルードしています。
9行目の module = torch::jit::load(“model to path/traced_model.pt”); ではPython側の処理で作成した変換された学習済みモデルを読み込んでいます。なお model to path は変換済みモデルが存在するパスに変更をしてください。
11行目の torch::Tensor input = torch::ones({1, 3, 256, 256}).to(“cuda”); では推論時に使用する入力サンプルデータを作成しています。実際に使用する際はPython側で指定したチャンネル数を合わせてお好みのデータを用意してください。ここでもPython側と同様に .to(“cuda”) と .to(“cpu”) は一致させる必要がありますのでご注意ください。
14行目の auto elements = module.forward(input).toTuple() -> elements(); では elements という変数に配列型で推論結果を格納しています。
16行目の auto output = elements[0].toTensor(); ではネットワークの出力結果の1番目を output という変数に格納しています。elements 変数にはニューラルネットからの推論結果が順番に格納されているため、複数の出力がある場合は配列の中から別の要素を参照してください。
おわりに
今回は「Pytorchで作成・学習したニューラルネットをC++で使用する方法」を紹介しました。
今回使用した方法はC++以外の言語やサービスにも使用することができるため、準備が整い次第また紹介していきたいと思います。
本ブログではこれからも、管理人が気になったことや学んだことを共有していきますので、ぜひフォローお願いします。