pointnet C++推理部署--libtorch框架
由于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")
python推理代码：
import numpy as np
import torch
import pointnet_cls as model
num_point = 1024
num_class = 10
def pc_normalize(pc):
centroid = np.mean(pc, axis=0)
pc = pc - centroid
m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
pc = pc / m
return pc
if __name__ == '__main__':
file = r'./bed_0610.txt'
data = np.loadtxt(file, delimiter=',').astype(np.float32)
point_set = data[:, 0:3]
point_set = point_set[0:num_point, :]
point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
points = torch.from_numpy(point_set)
points = torch.reshape(points,((1, num_point, 3)))
label = torch.tensor([[0]], dtype=torch.int32)
classifier = model.get_model(num_class, normal_channel=False)
classifier.load_state_dict(torch.load('./best_model.pth'))
with torch.no_grad():
classifier = classifier.eval()
points = points.transpose(2, 1)
pred, _ = classifier(points)
pred_choice = pred.data.max(1)[1]
pred_list = pred_choice.cpu().numpy().tolist()
print(pred_list)
C++推理代码：
#include 预测结果：
预测类别为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")
python推理代码：
import torch
import numpy as np
import pointnet_part_seg as model
num_point = 2048
num_classes = 1
num_part = 4
seg_classes = {'Airplane': [0, 1, 2, 3]}
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
return new_y
def pc_normalize(pc):
centroid = np.mean(pc, axis=0)
pc = pc - centroid
m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
pc = pc / m
return pc
if __name__ == '__main__':
file = r'85a15c26a6e9921ae008cc4902bfe3cd.txt'
data = np.loadtxt(file).astype(np.float32)
point_set = data[:, 0:3]
point_set[:, 0:3] = pc_normalize(point_set[:, 0:3])
choice = np.random.choice(point_set.shape[0], num_point, replace=True)
point_set = point_set[choice, :][:, 0:3]
pts = point_set
points = torch.from_numpy(point_set)
points = torch.reshape(points,((1, num_point, 3)))
label = torch.tensor([[0]], dtype=torch.int32)
classifier = model.get_model(num_part, normal_channel=False)
classifier.load_state_dict(torch.load('./best_model.pth'))
with torch.no_grad():
classifier = classifier.eval()
points, label, = points.float(), label.long()
cloud = points.cpu().data.numpy()
points = points.transpose(2, 1)
seg_pred, _ = classifier(points, to_categorical(label, num_classes))
cur_pred_val = seg_pred.cpu().data.numpy()
cur_pred_val_logits = cur_pred_val
cur_pred_val = np.zeros((1, num_point)).astype(np.int32)
logits = cur_pred_val_logits[0, :, :]
cur_pred_val[0, :] = np.argmax(logits, 1)
pts = np.append(cloud.reshape(num_point, 3), cur_pred_val[0, :].reshape(num_point, 1), 1)
np.savetxt('pred.txt', pts, fmt='%.06f')
C++推理代码：
#include 预测结果：
参考：Libtorch部署模型
在C+中部署python（libtoch）模型的方法总结+,PytorchLibtorch,Win10VS2017
A simple C++ implementation of Charles Qi’s PointNet
已经转换的pt文件下载地址：pointnet torchscipt转换得到的pt文件