Window C++ ONNX 환경 실행

지난 시간에 C++에서 ONNX 환경 테스트를 실행하였는데요,

 

이번 시간에는 PyTorch에서 딥러닝 모델을 C++ onnx로 변경하는 방법에 대해 써보겠습니다.

 

첫단계로, PyTorch 사이트에서 https://pytorch.org/docs/stable/onnx.html#id2에서 model을 onnx로 변경하는 방법을 숙지합니다.

 

torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)

 

보통은 학습이 된 model을 옮기실텐데, model 선언 이후, torch.load, model.load_state_dict로 모델에 weight값을 불러와야 합니다.

 

위의 단계를 하기 위해 dummy_input이라고, input type을 맞춘 임의의 np.array를 생성하시면 됩니다.

 

그리고 output_names 를 설정하셔야 하는데, 이유는 아래에서 알려드리겠습니다.

 

생성된 onnx 파일을 이전의 C++ 프로젝트에 옮기고, 마찬가지로 소스코드 내의 모델명.onnx 파일을 변경합니다.

 

추가로 GPU를 사용할 경우, 추가로 cuda_provider_factory.h 헤더 파일을 임포트하여

 

OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0); 을 추가 하시면 됩니다.

 

여기서 파라미터의 숫자는 GPU의 번호를 의미합니다.

 

넣은 예시는 다음과 같습니다.

int main(void) {
	Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
	Ort::SessionOptions session_options;
	session_options.SetThreadPoolSize(1);
	OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0); // <<GPU 추가한 경우,
    //0의 경우 0번 GPU를 사용한 예시이다.
	session_options.SetGraphOptimizationLevel(1);

#ifdef _WIN32
	const wchar_t* model_path = L"squeezenet1.1.onnx";
#else
	const char* model_path = "squeezenet1.1.onnx";
#endif

	printf("Using Onnxruntime C++ API\n");
	Ort::Session session(env, model_path, session_options);

	Ort::Allocator allocator = Ort::Allocator::CreateDefault();

	size_t num_input_nodes = session.GetInputCount();
	std::vector<const char*> input_node_names(num_input_nodes);
	std::vector<int64_t> input_node_dims;

	// iterate over all input nodes
	for (int i = 0; i < num_input_nodes; i++) {
		// print input node names
		char* input_name = session.GetInputName(i, allocator);
		input_node_names[i] = input_name;

		// print input node types
		Ort::TypeInfo type_info = session.GetInputTypeInfo(i);
		auto tensor_info = type_info.GetTensorTypeAndShapeInfo();

		// print input shapes/dims
		input_node_dims = tensor_info.GetShape();
	}

	size_t input_tensor_size = 224 * 224 * 3;

	std::vector<float> input_tensor_values(input_tensor_size);
	std::vector<const char*> output_node_names = { "squeezenet0_flatten0_reshape0" };

	for (unsigned int i = 0; i < input_tensor_size; i++)
		input_tensor_values[i] = (float)i / (input_tensor_size + 1);

	Ort::AllocatorInfo allocator_info = Ort::AllocatorInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
	Ort::Value input_tensor = Ort::Value::CreateTensor<float>(allocator_info, input_tensor_values.data(), input_tensor_size, input_node_dims.data(), 4);
	assert(input_tensor.IsTensor());

	auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_node_names.data(), &input_tensor, 1, output_node_names.data(), 1);
	assert(output_tensors.size() == 1 && output_tensors.front().IsTensor());

	float* floatarr = output_tensors.front().GetTensorMutableData<float>();

	for (int i = 0; i < 5; i++)
		printf("Score for class [%d] =  %f\n", i, floatarr[i]);

	printf("Done!\n");
	return 0;
}

위의 기존 코드에서 모델에 따라 바꿀 것이 여러가지 있는데, 위에서부터 하나씩 말씀드리면,

 

처음에 봐야될 될 것이 input_tensor_size 입니다.

size_t input_tensor_size = 224 * 224 * 3;

아무래도 C++은 메모리를 직접 할당해야 하는 구조이다 보니 입력 사이즈의 크기를 보는 것처럼 맞춰야 합니다.

 

제일 중요한 것은 C++에서는 다차원 배열이 지원하지 않기 때문에, Input의 차원을 1차원 Vector 형태로 바꿔야 합니다. 

 

저의 경우는 이미지를 다루는데 이미지는 RGB를 이용한 3차원이므로 특정 픽셀(x,y)의 R,G,B가 100,110,120 이였다면 [100, 110, 120] 으로 변경되서 적용시켜야 합니다.

 

다음으로 볼 것은 output_node_names 인데 여기서 위에 말했던 output_names로 설정하지 않을 경우 에러가 발생합니다.

std::vector<const char*> output_node_names = { "squeezenet0_flatten0_reshape0" };

또, 데이터의 사이즈에 맞춰서 input_tensor를 생성하게 되는데 이 때, 입력의 차원 수를 맞춰서 변경해야 합니다.

우측의 4를 보면 입력이 batch size를 포함하면 입력이 1 x 224 x 224 x 3 으로 4차원이 되기 때문에 4가 됩니다.

Ort::Value input_tensor = Ort::Value::CreateTensor<float>(allocator_info, input_tensor_values.data(), input_tensor_size, input_node_dims.data(), 4);

input_tensors도 맞추었으니 마찬가지로 출력인 output_tensors도 차원 수를 맞춰야 합니다. 마찬가지로 마지막 우측에

차원 수를 맞춰서 넣어주시면 됩니다.

auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_node_names.data(), &input_tensor, 1, output_node_names.data(), 1);

 

저의 경우에는 동영상을 주로 다루는데 동영상의 경우, W x H x 3 x L(프레임 길이) 에 batch를 포함하면 5차원이 됩니다.

 

별도의 전처리를 포함해서 실행하는 저의 코드는 다음과 같습니다.

 

#include <assert.h>
#include <vector>
#include <onnxruntime_cxx_api.h>
#include "providers/cuda/cuda_provider_factory.h"
#include "opencv/highgui.h"
#include "opencv2/opencv.hpp"
#include <iostream>

const int frame_length = 16;
std::vector<float> Vector_change(std::vector<cv::Mat> input);
int main(int argc, char* argv[]) {
	Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
	Ort::SessionOptions session_options;
	session_options.SetThreadPoolSize(1);
	OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
	session_options.SetGraphOptimizationLevel(1);

#ifdef _WIN32
	const wchar_t* model_path = L"R2plus1D_model_v3.onnx";
#else
	const char* model_path = "R2plus1D_model.onnx";
#endif

	printf("Using Onnxruntime C++ API\n");
	Ort::Session session(env, model_path, session_options);

	Ort::Allocator allocator = Ort::Allocator::CreateDefault();
	size_t num_input_nodes = session.GetInputCount();
	std::vector<const char*> input_node_names(num_input_nodes);
	input_node_names.reserve(session.GetInputCount());
	std::vector<int64_t> input_node_dims;

	// iterate over all input nodes
	for (int i = 0; i < num_input_nodes; i++) {
		// print input node names
		char* input_name = session.GetInputName(i, allocator);
		input_node_names[i] = input_name;

		// print input node types
		Ort::TypeInfo type_info = session.GetInputTypeInfo(i);
		auto tensor_info = type_info.GetTensorTypeAndShapeInfo();

		// print input shapes/dims
		input_node_dims = tensor_info.GetShape();
	}

	size_t input_tensor_size = 1 * 3 * 16 * 224 * 224;


	cv::VideoCapture cap("Video.mp4");
	std::vector<cv::Mat> imgs;
	while (cap.isOpened()) {
		cv::Mat frame;
		cap >> frame;
		frame.convertTo(frame, CV_32FC3);
		frame /= (float)255.0;
		
		cv::resize(frame, frame, cv::Size(224, 224));
		cv::imshow("a1", frame);
		cv::waitKey(1);
		
		if(imgs.size() < 16)
			imgs.push_back(frame);
		else {
			imgs.push_back(frame);
			imgs.erase(imgs.begin(), imgs.begin()+1);
			std::vector<float> input_tensor_values = Vector_change(imgs);
			/*for (unsigned int i = 0; i < input_tensor_size; i++)
				input_tensor_values[i] = (float)i / (input_tensor_size + 1);*/
			std::vector<const char*> output_node_names = { "output1" };
			Ort::AllocatorInfo allocator_info = Ort::AllocatorInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
			Ort::Value input_tensor = Ort::Value::CreateTensor<float>(allocator_info, input_tensor_values.data(), input_tensor_size, input_node_dims.data(), 5);
			assert(input_tensor.IsTensor());
			auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_node_names.data(), &input_tensor, 1, output_node_names.data(), 1);
			assert(output_tensors.size() == 1 && output_tensors.front().IsTensor());
			float* floatarr = output_tensors.front().GetTensorMutableData<float>();
			float output = 1 / (1 + std::exp(-*floatarr));
			printf("Score for class =  %f %f\n", *floatarr, output);
		}
	}
	cap.release();
	return 0;
	
}
std::vector<float> Vector_change(std::vector<cv::Mat> imgs) {
	std::vector<float> input_tensor_values;
	for (int i = 0; i < frame_length; i++) {
		for (int j = 0; j < imgs[i].rows; j++) {
			input_tensor_values.insert(input_tensor_values.end(), imgs[i].ptr<float>(j), imgs[i].ptr<float>(j) + imgs[i].cols * 3);
		}
	}
	return input_tensor_values; 
}