ONNX Runtimeを使ってみる - TadaoYamaokaの開発日記

dlshogiはCUDAに対応したNvidiaのGPUが必須になっているが、AMDのGPUやCPUのみでも動かせるようにしたいと思っている。

Microsoftがオープンソースで公開しているONNX Runtimeを使うと、様々なデバイスでONNXモデルの推論を行うことができる。
TensorRT対応で、ONNXのモデルを読み込めるようになったので、ONNX Runtimeに対応すれば同じモデルを使いまわせる。

ONNX Runtimeは、PythonやC#など複数の言語のインターフェースが提供されている。
dlshogiに組み込むにはC++のインターフェースが必要だが、C++も提供されている。

推論に使うデバイスは、CPUやCUDA、TensorRT、DirectX、MKL-DNNなど複数のデバイスを切り替えられるようになっている。
DirectXに対応すれば、AMDのGPUでも高速に推論が可能になる。

ということで、練習のためにONNX RuntimeをC++から使って、MNISTの推論を行ってみた。
今回は、Linux(Ubuntu 18.04)でCPUのみで動かした。

ONNX Runtimeのインストール

GitHubのReleaseからバイナリをダウンロードする。
Linux向けに、onnxruntime-linux-x64-1.3.0.tgzをダウンロードし、適当なディレクトリに展開する。

LD_LIBRARY_PATHの設定

export LD_LIBRARY_PATH=/xxxx/onnxruntime-linux-x64-1.3.0/lib:$LD_LIBRARY_PATH

サンプルコード

公式のリポジトリで、C++のMNISTの推論のサンプルが提供されている。
しかし、これはWindowsのWin32APIを使ったGUIプログラムになっている。
（クロスプラットフォームのオープンソースのサンプルコードがWin32プログラムとかやめてほしい・・・）

Win32のコードを削除して、コンソールプログラムにした。

#include <onnxruntime_cxx_api.h>
#include <algorithm>
#include <iostream>

struct MNIST {
	MNIST() {
		auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
		input_tensor_ = Ort::Value::CreateTensor<float>(memory_info, input_image_.data(), input_image_.size(), input_shape_.data(), input_shape_.size());
		output_tensor_ = Ort::Value::CreateTensor<float>(memory_info, results_.data(), results_.size(), output_shape_.data(), output_shape_.size());
	}

	std::ptrdiff_t Run() {
		const char* input_names[] = {"Input3"};
		const char* output_names[] = {"Plus214_Output_0"};

		session_.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor_, 1, output_names, &output_tensor_, 1);

		result_ = std::distance(results_.begin(), std::max_element(results_.begin(), results_.end()));
		return result_;
	}

	static constexpr const int width_ = 28;
	static constexpr const int height_ = 28;

	std::array<float, width_ * height_> input_image_{};
	std::array<float, 10> results_{};
	int64_t result_{0};

private:
	Ort::Env env;
	Ort::Session session_{env, "model.onnx", Ort::SessionOptions{nullptr}};

	Ort::Value input_tensor_{nullptr};
	std::array<int64_t, 4> input_shape_{1, 1, width_, height_};

	Ort::Value output_tensor_{nullptr};
	std::array<int64_t, 2> output_shape_{1, 10};
};

std::unique_ptr<MNIST> mnist_;

int main()
{
	try {
		mnist_ = std::make_unique<MNIST>();
	} catch (const Ort::Exception& exception) {
		std::cerr << exception.what() << std::endl;
		return 1;
	}

	mnist_->input_image_ = {
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
		0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
	};
	mnist_->Run();
	std::cout << mnist_->result_ << std::endl;

	return 0;
}

Makefile

ビルドのためにMakefileを作成する。

CC = g++
CFLAGS = -std=c++14 -msse4.2 -mbmi2
LDFLAGS = -lonnxruntime
INCLUDE = -I../onnxruntime-linux-x64-1.3.0/include
LIB = -L../onnxruntime-linux-x64-1.3.0/lib

target = mnist
sources = mnist.cpp
objects = $(addprefix obj/, $(sources:.cpp=.o))

$(target): $(objects)
	$(CC) -o $@ $^ $(LIB) $(LDFLAGS) $(CFLAGS)

obj/%.o: %.cpp
	@[ -d obj ] || mkdir -p obj
	$(CC) $(CFLAGS) $(INCLUDE) -o $@ -c $<

all: $(target)

clean:
	rm -f $(objects) $(target)

ビルド

make

モデルダウンロード

サンプルの説明にあるMNISTのONNXモデルをダウンロードする。
https://github.com/onnx/models/tree/master/vision/classification/mnist

ONNX version 1.3のモデルをダウンロードした。
展開して、model.onnxを実行ファイルと同じディレクトリに配置する。

実行

./mnist

結果：

まとめ

ONNX RuntimeをC++から使用してMNISTの推論を行うことができた。

ONNX Runtimeは計算グラフの最適化を行う機能があるが、今回のサンプルでは試していない。
GPUと比較した場合のパフォーマンスも気になるところなので、別途比較してみたい。