Pointnet C++ Win10部署

一、环境配置：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快了很多。

PartSegmentation
pytorch训练得到的pth文件转libtorch使用的pt文件脚本（以1类物体分成4部分，gpu版本为例）：
torchscript.py

import torch
import pointnet_part_seg

def to_categorical(y, num_classes):
    """ 1-hot encodes a tensor """
    new_y = torch.eye(num_classes)[y.cpu().data.numpy(),]
    if (y.is_cuda):
        return new_y.cuda()
    return new_y

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

example=torch.rand(1, 3, 2048)
example=example.cuda() #cpu版本需注释此句
label=torch.rand(1, 1)
label=label.cuda() #cpu版本需注释此句

traced_script_module = torch.jit.trace(model, (example, to_categorical(label, 1)))
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 resample(std::vector<float>& points, int nums)
{
	srand((int)time(0));
	std::vector<int> choice(nums);
	for (size_t i = 0; i < nums; i++)
	{
		choice[i] = rand() % (points.size() / 3);
	}

	std::vector<float> temp_points(3 * nums);
	for (size_t i = 0; i < nums; i++)
	{
		temp_points[3 * i] = points[3 * choice[i]];
		temp_points[3 * i + 1] = points[3 * choice[i] + 1];
		temp_points[3 * i + 2] = points[3 * choice[i] + 2];
	}
	points = temp_points;
}

at::Tensor classfier(std::vector<float>& points, std::vector<float>& labels)
{
	torch::Tensor points_tensor = torch::from_blob(points.data(), { 1, 2048, 3 }, torch::kFloat);
	torch::Tensor labels_tensor = torch::from_blob(labels.data(), { 1, 1, 1}, torch::kFloat);

	points_tensor = points_tensor.to(torch::kCUDA);
	points_tensor = points_tensor.permute({ 0, 2, 1 });
	//std::cout << points_tensor << std::endl;
	labels_tensor = labels_tensor.to(torch::kCUDA);
	//std::cout << labels_tensor << std::endl;

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

	auto outputs = module.forward({ points_tensor, labels_tensor }).toTuple();
	torch::Tensor out0 = outputs->elements()[0].toTensor();
	//std::cout << out0 << std::endl; //[ CUDAFloatType{1,2048,4} ]
	out0 = torch::squeeze(out0);
	//std::cout << out0 << std::endl; //[ CUDAFloatType{2048,4} ]

	auto max_classes = out0.max(1);
	auto max_result = std::get<0>(max_classes);
	auto max_index = std::get<1>(max_classes);
	//std::cout << max_result << std::endl;
	//std::cout << max_index << std::endl;
	
	return max_index;
}


int main()
{
	std::vector<float> points, labels;
	float x, y, z, nx, ny, nz, label;
	int point_num = 2048;
	std::ifstream infile;
	infile.open("85a15c26a6e9921ae008cc4902bfe3cd.txt");
	while (infile >> x >> y >> z >> nx >> ny >> nz >>label)
	{
		points.push_back(x);
		points.push_back(y);
		points.push_back(z);
	}
	labels.push_back(1.0);
	infile.close();

	pc_normalize(points);
	resample(points, point_num);

	at::Tensor result = classfier(points, labels);

	std::fstream outfile;
	outfile.open("85a15+.txt", 'w');
	for (size_t i = 0; i < point_num; i++)
	{
		outfile << points[3 * i] << " " << points[3 * i + 1] << " " << points[3 * i + 2]  << " " << result[i].item<int>() << std::endl;
	}
	outfile.close();

	system("pause");
	return 0;
}

预测结果：

参考文献：
https://blog.csdn.net/taifyang/article/details/124257666
https://blog.csdn.net/taifyang/article/details/124332344

扫描二维码关注公众号，回复： 14347104 查看本文章