由于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