Libtorch resnet baseline 코드 (trace version)

Libtorch를 사용하는 방법은 크게 2가지로 나뉩니다.

1. pytorch으로 학습 후 weight를 torchscript로 변환 후 libtorch로 load 하기

2. libtorch로 학습 후 그대로 load 하기

 

이 중 대부분은 1번으로 libtorch를 사용합니다. 왜냐하면 오픈 소스가 대부분 python으로 작성돼있기 때문입니다.

만약 본인이 실력이 있고 논문을 구현할 수 있을 정도라면 2번을 추천합니다. (하지만 python이나 c++이나 딥러닝 부분에서는 큰 차이가 없습니다.)

 

1번에서 torchscript 변환은 2가지 방법으로 진행할 수 있습니다.

(1) tracing을 통한 변환

(2) annotation을 통한 변환

tracing을 통한 변환은 defualt 변환이라고 생각하시면 되고, 만약 model class 코드 forward method에 if문과 같이 제어 흐름이 존재한다면 annotation을 사용하면 됩니다.

 

오늘은 cpu 버전 tracing을 기준으로 진행하겠습니다.

 

깃헙 링크: https://github.com/yunseokddi/LibTorchDev/tree/master/Libtorch%20Baseline

yunseokddi/LibTorchDev

 

libtorch를 진행하기 전 위에서 언급한 방법과 같이 우선 python으로 학습을 진행한 후 trace 변환을 진행해야 합니다.

import torch
import torchvision

# Download pretrained weight
model = torchvision.models.resnet18(pretrained=True)

# Don't forget change model to eval mode
model.eval()
example = torch.rand(1, 3, 224, 224)

# trace weight
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("./traced_resnet18.pt")

코드를 살펴보자면 우선 간단하게 resnet18 모델을 pretrain을 하고 불러옵니다.

그리고 여기서 꼭 기억해야 할 사실이 있습니다. 반드시 eval()을 선언해야합니다. (이것 때문에 몇 시간 고생했습니다.)

example 변수를 선언하여 입력 데이터의 shape를 지정합니다.

그 후 torch.jit.trace를 통해 torchscript 변환을 하여 save로 저장을 합니다.

경로를 확인해보면 torchscript로 변환된 weight 파일을 확인하실 수 있습니다.

 

우선 전체 libtorch 코드를 보고 살펴보도록 하겠습니다.

libtorch 비주얼 스튜디오 설치 방법 링크: https://ys-cs17.tistory.com/43

Visual Studio 2019에서 LibTorch 사용하기

#include <torch/script.h>
#include <opencv2/opencv.hpp>

#include <iostream>
#include <memory>
#include <string>
#include <vector>

#define IMAGE_SIZE 224
#define CHANNELS 3

using namespace std;

bool loadimage(std::string file_name, cv::Mat& image)
{
	image = cv::imread(file_name);

	cv::cvtColor(image, image, cv::COLOR_BGR2RGB);

	// scale image to fit
	cv::Size scale(IMAGE_SIZE, IMAGE_SIZE);
	cv::resize(image, image, scale);
	// convert [unsigned int] to [float]
	image.convertTo(image, CV_32FC3, 1.0f / 255.0f);

	return true;
}

bool loadimagenetlabel(std::string file_name, std::vector<std::string>& labels)
{
	std::ifstream ifs(file_name);
	if (!ifs)
	{
		return false;
	}
	std::string line;
	while (std::getline(ifs, line))
	{
		labels.push_back(line);
	}
	return true;
}

int main() {
	// change your weight, label, img paths
	const string weight_path = "../weight/traced_resnet18.pt";
	const string label_path = "../data/label.txt";
	const string img_path = "../data/sample_image/dog.jpg";

	// load the pretrained model
	torch::jit::script::Module module = torch::jit::load(weight_path);

	std::vector<std::string> labels;

	// mapping from label.txt to idx
	loadimagenetlabel(label_path, labels);

	cv::Mat image;

	if (loadimage(img_path, image))
	{
		// mat to tensor
		auto input_tensor = torch::from_blob(image.data, { 1, IMAGE_SIZE, IMAGE_SIZE, CHANNELS });
		input_tensor = input_tensor.permute({ 0, 3, 1, 2 });

		// model forward
		torch::Tensor out_tensor = module.forward({ input_tensor }).toTensor();

		auto results = out_tensor.sort(-1, true);
		auto softmaxs = std::get<0>(results)[0].softmax(0);
		auto indexs = std::get<1>(results)[0];


		auto idx = indexs[0].item<int>();
		std::cout << "label:  " << labels[idx] << std::endl;
		std::cout << "probability:  " << softmaxs[0].item<float>() * 100.0f << "%" << std::endl;



	}
	return 0;
}

Global 부분부터 살펴보면 우선 libtorch와 opencv를 incude해준후 image size와 channel을 224, 3으로 define 해줍니다.

bool loadimage(std::string file_name, cv::Mat& image)
{
	image = cv::imread(file_name);

	cv::cvtColor(image, image, cv::COLOR_BGR2RGB);

	// scale image to fit
	cv::Size scale(IMAGE_SIZE, IMAGE_SIZE);
	cv::resize(image, image, scale);
	// convert [unsigned int] to [float]
	image.convertTo(image, CV_32FC3, 1.0f / 255.0f);

	return true;
}

다음은 loadimage 함수입니다. 말 그대로 이미지를 읽고 preprocessing을 진행한 후 call by reference로 변환해주는 함수입니다. opencv는 이미지를 BGR로 읽기 때문에 이를 RGB로 변환해 주어야 합니다. 그러고 나서 이미지를 model input size와 맞게 resize 한 후 float 형태로 변환시킵니다.

bool loadimagenetlabel(std::string file_name, std::vector<std::string>& labels)
{
	std::ifstream ifs(file_name);
	if (!ifs)
	{
		return false;
	}
	std::string line;
	while (std::getline(ifs, line))
	{
		labels.push_back(line);
	}
	return true;
}

해당 함수는 label과 index를 mapping 하기 위해 제작했습니다. pytorch resnet은 imagenet을 통해 학습한 모델이라 lmagenet 라벨이 필요합니다.

labels 변수에 label.txt 파일을 한 줄씩 읽어 push 하는 역할을 합니다.

	const string weight_path = "../weight/traced_resnet18.pt";
	const string label_path = "../data/label.txt";
	const string img_path = "../data/sample_image/dog.jpg";

	// load the pretrained model
	torch::jit::script::Module module = torch::jit::load(weight_path);

	std::vector<std::string> labels;

	// mapping from label.txt to idx
	loadimagenetlabel(label_path, labels);

다음으로는 main 부분을 살펴봅시다.

각 weight, label, image path를 지정해줍니다. 그리고 module을 선언하여 weight를 load 해줍니다.

torchscript로 변환된 weight를 읽을 때는 따로 model을 제작하지 않으셔도 됩니다. 

그 이유는 앞서 python에서 torchscript 변환할 때 model의 정보까지 trace 해 모델 정보까지 저장되어 있기 때문입니다.

그러고 나서 vector의 형태로 labels를 선언해 주고, 이를 loadimagenetlabel 함수로 넣어줍니다.

cv::Mat image;

if (loadimage(img_path, image))
	{
		// mat to tensor
		auto input_tensor = torch::from_blob(image.data, { 1, IMAGE_SIZE, IMAGE_SIZE, CHANNELS });
		input_tensor = input_tensor.permute({ 0, 3, 1, 2 });

		// model forward
		torch::Tensor out_tensor = module.forward({ input_tensor }).toTensor();

		auto results = out_tensor.sort(-1, true);
		auto softmaxs = std::get<0>(results)[0].softmax(0);
		auto indexs = std::get<1>(results)[0];


		auto idx = indexs[0].item<int>();
		std::cout << "label:  " << labels[idx] << std::endl;
		std::cout << "probability:  " << softmaxs[0].item<float>() * 100.0f << "%" << std::endl;
	}
	return 0;

이후 Mat 자료형인 image 변수를 생성하고, loadimage를 통해 preprocess 과정이 완료된 이미지를 읽습니다.

읽은 이미지를 from_blob을 통해 mat to tensor로 변환시키면서 이미지의 차원을 1개 더 늘립니다. 그리고 나서 permute를 통해 기존 [batch, width, height, channels]를 [batch, channels, width, channels]로 변환시킵니다.

이렇게 최종적으로 변환된 이미지를 module.forward로 model에 넣어주고 result를 얻습니다.

해당 result를 정렬화시켜 softmax를 진행한 후 가장 높은 확률과 이에 mapping 된 index 값을 얻습니다.

그러고 나서 최종적으로 해당 index에 해당되는 label 값과 확률 값을 얻게 됩니다.

output

다음은 해당 libtorch 코드와 똑같은 python 코드를 제작하여 결괏값을 비교해봅시다. 코드 설명은 생략하겠습니다.

import torch
import torchvision
import cv2

from torchvision import transforms

data_transform = transforms.Compose([
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                              std=[0.229, 0.224, 0.225]),
    transforms.ToTensor()
])

def preprocess(img):
    # scale image to fit
    img = cv2.resize(img, (224, 224))
    img = img[:, :, ::-1].transpose((2, 0, 1)).copy()
    # convert [unsigned int] to [float]
    img = torch.from_numpy(img).float().div(255.0).unsqueeze(0)
    return img

def detect_img(img_path):
    img = cv2.imread(img_path)
    img = preprocess(img)

    with torch.no_grad():
        output = model(img)
        predict = torch.argmax(output)
    return predict.item(), max(torch.softmax(max(output), dim=0)).item()

def idx_to_label(label_path, idx):
    label_file = open(label_path, 'r', encoding='UTF-8')
    lines = label_file.readlines()
    return lines[idx]


if __name__ == "__main__":
    img_path = '../sample_data/dog.jpg'
    label_path = '../label.txt'

    # use pretrained model
    model = torchvision.models.resnet18(pretrained=True)
    model.eval()

    pred_idx, prob = detect_img(img_path)

    pred_label = idx_to_label(label_path, pred_idx)

    print("Label: {}".format(pred_label))
    print("Probability: {}%".format(round(prob*100, 4)))

output

최종적으로 같은 output이 나오는 것을 확인할 수 있습니다.