pointnet C++推理部署(一)

由于tensorflow编译C++的api比较麻烦,此次部署的pointnet代码的Python版本为Pytorch编写的。
代码:Pointnet_Pointnet2_pytorch
环境配置:win10系统
cuda10.1+cudnn7.5+Python3.6.5+pytorch1.5.0+libtorch1.5.0+VS2017
或者libtorch1.4.0+VS2015
软件下载和配置过程在此不赘述。

classification

pytorch训练得到的pth文件转libtorch使用的pt文件脚本(以分10类,gpu版本为例):
torchscript.py

import torch
import pointnet_cls

model = pointnet_cls.get_model(10, False)
model = model.cuda() #cpu版本需注释此句
model.eval()
model.load_state_dict(torch.load('best_model.pth'))

example=torch.rand(1, 3, 1024)
example=example.cuda() #cpu版本需注释此句
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("best_model.pt")

C++部署代码:

#include <iostream>
#include <vector>
#include <fstream>
#include <torch/script.h>


void pc_normalize(std::vector<float>& points)
{
    
    
	int N = points.size() / 3;
	float mean_x = 0, mean_y = 0, mean_z = 0;
	for (size_t i = 0; i < N; ++i)
	{
    
    
		mean_x += points[3 * i];
		mean_y += points[3 * i + 1];
		mean_z += points[3 * i + 2];
	}
	mean_x /= N;
	mean_y /= N;
	mean_z /= N;

	for (size_t i = 0; i < N; ++i)
	{
    
    
		points[3 * i] -= mean_x;
		points[3 * i + 1] -= mean_y;
		points[3 * i + 2] -= mean_z;
	}

	float m = 0;
	for (size_t i = 0; i < N; ++i)
	{
    
    
		if (sqrt(pow(points[3 * i], 2) + pow(points[3 * i + 1], 2) + pow(points[3 * i + 2], 2)) > m)
			m = sqrt(pow(points[3 * i], 2) + pow(points[3 * i + 1], 2) + pow(points[3 * i + 2], 2));
	}

	for (size_t i = 0; i < N; ++i)
	{
    
    
		points[3 * i] /= m;
		points[3 * i + 1] /= m;
		points[3 * i + 2] /= m;
	}
}


void classfier(std::vector<float>& points)
{
    
    
	torch::Tensor points_tensor = torch::from_blob(points.data(), {
    
     1, 1024, 3 }, torch::kFloat);
	points_tensor = points_tensor.to(torch::kCUDA);
	points_tensor = points_tensor.permute({
    
     0, 2, 1 });
	//std::cout << points_tensor << std::endl;

	torch::jit::script::Module module = torch::jit::load("classes10_gpu.pt");
	module.to(torch::kCUDA);

	auto outputs = module.forward({
    
    points_tensor}).toTuple();
	torch::Tensor out0 = outputs->elements()[0].toTensor();
	std::cout << out0 << std::endl;

	auto max_result = out0.max(1, true);
	auto max_index = std::get<1>(max_result).item<int>();
	std::cout << max_index << std::endl;
}


int main()
{
    
    
	std::vector<float> points;
	std::ifstream infile;
	float x, y, z, nx, ny, nz;
	char ch;
	infile.open("bed_0610.txt");
	int point_num = 0;
	while (infile >> x >> ch >> y >> ch >> z >> ch >> nx >> ch >> ny >> ch >> nz)
	{
    
    
		points.push_back(x);
		points.push_back(y);
		points.push_back(z);
		++point_num;
		if (point_num == 1024)	break;
	}
	infile.close();

	pc_normalize(points);

	classfier(points);

	system("pause");
	return 0;
}

预测结果:
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述
预测类别为1,在names.txt中对应为bed,结果正确。
C++推理速度稳定在不到0.2s,相比Python推理速度1~2s快了很多。
参考:Libtorch部署模型
在C+中部署python(libtoch)模型的方法总结+,PytorchLibtorch,Win10VS2017
A simple C++ implementation of Charles Qi’s PointNet

猜你喜欢

转载自blog.csdn.net/taifyang/article/details/124257666#comments_22029726
今日推荐