落叶无痕
0%

LIbtorch与pytorch相似与不同

发表于 分类于 Changyan：

自动摘要: 💡根据遗忘曲线： ……..

💡 根据 遗忘曲线：笔记正是帮助你记录和回顾的工具，不必拘泥于形式，其核心是：记录、翻看、思考

主题 如何使用libtorch实现深度学习
状态 基础版完成
简介 目前，国内各大平台似乎没有pytorch在c++上api的完整教学，也没有基于c++开发的完整的深度学习开源模型。可能原因很多：1.** c/c++的深度学习已经足够底层和落地，商用价值较高，开发难度偏大，一般不会开源**； 2. 基于python训练，libtorch预测的部署形式足够满足大多数项目的需求，如非产品级应用，不会有人愿意研究如何用c++从头搭建模型，实现模型训练功能； 3. Tensorflow的市场份额，尤其时工业应用的部署下市场占比足够高，导致基于libtorch的开发和部署占比很小，所以未见开源。

本节主要思路图

3D数据分割–环境配置

数据准备

数据制作，分为3D点云数据，对应标签数据


Cmake配置

用于配置基础编译环境

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
cmake_minimum_required(VERSION 3.21)
project(libtorch_pointnet)
message(STATUS "The current system is:${CMAKE_SYSTEM_NAME}  ")

# 配置C++变量
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_POLICY_DEFAULT_CMP0054 NEW)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
set(Open3D_DIR D:/software/develop/CompileLibrary/open3d-0.15.1/CMake)
set(Torch_DIR  D:/software/develop/CompileLibrary/libtorch/share/cmake/Torch)
set(EXECUTABLE_OUTPUT_PATH "${CMAKE_CURRENT_SOURCE_DIR}/bin/")
message(STATUS "OUTPUT_PATH: ${EXECUTABLE_OUTPUT_PATH}")
Set(MKLDNN_CPU_RUNTIME STREQUAL "TBB")
#SET(MKLDNN_THREADING "TBB")


# 参数
option(USE_SYSTEM_LIBS "Use all available system-provided libraries." OFF)
option(USE_CUDA "Use CUDA" OFF)

# 加载包
find_package(OpenCV 4.5.3 REQUIRED)
find_package(Open3D REQUIRED)
find_package(Torch REQUIRED)
find_package(fmt CONFIG REQUIRED)#用于fmt::format，fmt::print



# 编译为动态库
#add_library(${PROJECT_NAME} SHARED library.cpp)

# 编译为可执行程序
add_executable(${PROJECT_NAME}    main.cpp src/dataset.cpp include/dataset.h src/model.cpp include/model.h include/train.h src/train.cpp)

# 链接依赖
target_link_libraries(${PROJECT_NAME}  Open3D::Open3D ${TORCH_LIBRARIES} ${OpenCV_LIBS} fmt::fmt  )
set_property(TARGET ${PROJECT_NAME}  PROPERTY CXX_STANDARD 14)
# 拷贝文件
# 对于windows，需要加入以下选项
# https://github.com/pytorch/pytorch/issues/25457
if (MSVC)
  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
  add_custom_command(TARGET ${PROJECT_NAME}
  POST_BUILD
  COMMAND ${CMAKE_COMMAND} -E copy_if_different
  ${TORCH_DLLS} "D:\\software\\develop\\CompileLibrary\\open3d-0.15.1\\bin\\open3d.dll"
  $<TARGET_FILE_DIR:${PROJECT_NAME}>)
endif (MSVC)

if(USE_CUDA AND (CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 10.2) AND (CMAKE_HOST_SYSTEM_NAME MATCHES "Windows"))
# CUDA < 10.2 不支持在一次调用中编译和提取头依赖项，因此 CMake 两次调用 nvcc 并在其间使用 &&。
# 但是，在 Windows 上，cmd.exe 对我们开始点击的命令有 8191 个字符的限制。这会将大多数参数移动到文件中以避免超出限制
  set(CMAKE_CUDA_USE_RESPONSE_FILE_FOR_OBJECTS ON)
  set(CMAKE_NINJA_FORCE_RESPONSE_FILE ON CACHE INTERNAL "")
endif()


if(USE_SYSTEM_LIBS)
  set(USE_SYSTEM_CPUINFO ON)
  set(USE_SYSTEM_SLEEF ON)
  set(USE_SYSTEM_GLOO ON)
  set(BUILD_CUSTOM_PROTOBUF OFF)
  set(USE_SYSTEM_EIGEN_INSTALL ON)
  set(USE_SYSTEM_FP16 ON)
  set(USE_SYSTEM_PTHREADPOOL ON)
  set(USE_SYSTEM_PSIMD ON)
  set(USE_SYSTEM_FXDIV ON)
  set(USE_SYSTEM_BENCHMARK ON)
  set(USE_SYSTEM_ONNX ON)
  set(USE_SYSTEM_XNNPACK ON)
  set(USE_SYSTEM_PYBIND11 ON)

endif()


# 构建方式
# mkdir build
# cd build
# cmake  ..
# cmake --build . --config Release

自定义数据处理

按照libtorch规定格式进行加载

  1. 要求初始化类继承torch::data::Dataset
  2. 重写size()函数，用于返回数据大小
  3. 重写get(size_t index) 函数，用于数据加载器索引

C++版：

dataset.h

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
//
// Created by sindre on 2022/1/10.
// 自定义数据集
//

#ifndef LIBTORCH_TEMPLATE_DATASET_H
#define LIBTORCH_TEMPLATE_DATASET_H

#include <open3d/Open3D.h>
#include <torch/torch.h>
#include "opencv2/opencv.hpp"
#include "fmt/core.h"
#include <vector>
#include <string>
#include <iostream>
#include "torch/script.h"

using namespace torch::indexing;
class dataset : public torch::data::Dataset<dataset>{
    std::vector<std::string> data_path;
    std::vector<std::string> label_path;
    int npoints=2500;

    public:
        dataset(const std::string& root,size_t _npoints);
        //返回数据集大小
        torch::optional<size_t> size() const override;
        //返回数据内容
        torch::data::Example<> get(size_t index) override;

};


#endif//LIBTORCH_TEMPLATE_DATASET_H

dataset.cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
//
// Created by sindre on 2022/1/10.
//

#include "../include/dataset.h"

torch::data::Example<> dataset::get(size_t index) {

    //初始化，加载点云

    std::string _pcd_path = data_path[index];
    //std::cout<<_pcd_path<<std::endl;
    std::string _label_path = label_path[index];
    //fmt::print("In dataset get index pcd:{}",_pcd_path );
    std::vector<int> _label;
    //处理标签
    std::ifstream inFile(_label_path);
    if (!inFile.is_open())
    {
        //fmt::print("file read failed--{}", _label_path);
    }
    else {
        std::string line;
        while(getline(inFile,line)){
            _label.push_back(std::stoi(line));
        }
        inFile.close();
    }


    auto pcd=std::make_shared<open3d::geometry::PointCloud>();
    open3d::io::ReadPointCloud(_pcd_path,*pcd,{"xyz"});
    //随机采样
    std::vector<size_t> indices(pcd->points_.size());
    std::iota(std::begin(indices), std::end(indices), (size_t)0);
    //随机产生个随机数
    std::random_device rd;
    std::mt19937 rng(rd());
    std::shuffle(indices.begin(), indices.end(),rng);
    //固定点数
    std::vector<size_t> n_indices;
    for(int i =0 ;i<npoints;i++){
        n_indices.push_back(indices[i]);

    }
    pcd=pcd->SelectByIndex(n_indices);
    //fmt::print("SelectByIndex test {}", n_indices);
    //纠正label
    std::vector<int> _label_true;
    for (size_t i :n_indices){
        //std::cout<<_label[i]<<std::endl;
        _label_true.push_back(_label[i]);
    }
    //fmt::print("\n n_indices[0]--{} \n n_indices.size--{} \n _label_true[0]-->{} \n _label_true.size-->{}\n", n_indices[0],n_indices.size(),_label_true[0],_label_true.size());

    auto pcd_3d = pcd->points_;
    //std::cout << pcd_3d << std::endl;
    //std::cout << pcd_3d[1] << std::endl;
    torch::Tensor pcd_tensor = torch::zeros({3,npoints});  // my torch tensor for the forward pass
    for (int i = 0; i < npoints; i++) {
        pcd_tensor[0][i] = pcd_3d[i][0];
        pcd_tensor[1][i] = pcd_3d[i][1];
        pcd_tensor[2][i] = pcd_3d[i][2];
    }

    auto label_tensor = torch::tensor(_label_true);
    //fmt::print("\n label_tensor size: {}" ,label_tensor.sizes());
    //auto pcd_tensor  =torch::from_blob(pcd_3d.data(),{2500,3}).clone();

    //std::cout << pcd_tensor.sizes()<< std::endl;
    //fmt::print("\n pcd_tensor size:  {} {} \t {}" ,pcd_3d[0],pcd_3d.size(),pcd_tensor.sizes());
    //std::cout << pcd_tensor[0]<< std::endl;

    return {pcd_tensor,label_tensor};

    //pcd = pcd->UniformDownSample(dataset::options.npoints/pcd->points_.size());


    //缩放到[-1,1]
/*

    auto pcd_3d = open3d::core::eigen_converter::EigenMatrixToTensor(pcd->GetCenter()).ToFlatVector<float>();
    auto pcd_tensor=torch::from_blob(pcd_3d.data(),{int64_t(pcd_3d.size())},torch::TensorOptions().dtype(torch::kFloat32));
    auto pcd_center_tensor = pcd_tensor - torch::unsqueeze(torch::mean(pcd_tensor,0),0);
    auto dist = torch::amax(torch::sqrt(torch::sum(torch::pow(pcd_center_tensor,2),1)),0);
    auto point_set = torch::div(pcd_center_tensor , dist);
*/




}
torch::optional<size_t> dataset::size() const {
    return label_path.size();
}
dataset::dataset(const std::string& root, size_t _npoints) {
    //只针对现有数据集
    dataset::npoints = _npoints;
    std::vector<std::string> _data_path;
    cv::glob(root,_data_path,false);
    //fmt::print("dataset start load file! \n");
    for( std::string i : _data_path)
    {
        //std::string _label_path = i.replace(i.rfind("pts"),i.size(),"seg");
        if (i.find("pts"  )!= std::string::npos ){
            //fmt::print("pcd-->{} \n",i);
            data_path.push_back(i);
        }else{
            //fmt::print("label-->{} \n",i);
            label_path.push_back(i);
        }




    }
    //fmt::print("test dataset data_path[0]-->{} \n",data_path[0]);
}


python版

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class ShapeNetDataset(data.Dataset):
    def __init__(self,
                 root,
                 npoints=2500,
                 classification=False,
                 class_choice=None,
                 split='train',
                 data_augmentation=True):
        self.npoints = npoints
        self.root = root
        self.catfile = os.path.join(self.root, 'synsetoffset2category.txt')
        self.cat = {}
        self.data_augmentation = data_augmentation
        self.classification = classification
        self.seg_classes = {}
        
        with open(self.catfile, 'r') as f:
            for line in f:
                ls = line.strip().split()
                self.cat[ls[0]] = ls[1]
        #print(self.cat)
        if not class_choice is None:
            self.cat = {k: v for k, v in self.cat.items() if k in class_choice}

        self.id2cat = {v: k for k, v in self.cat.items()}

        self.meta = {}
        splitfile = os.path.join(self.root, 'train_test_split', 'shuffled_{}_file_list.json'.format(split))
        #from IPython import embed; embed()
        filelist = json.load(open(splitfile, 'r'))
        for item in self.cat:
            self.meta[item] = []

        for file in filelist:
            _, category, uuid = file.split('/')
            if category in self.cat.values():
                self.meta[self.id2cat[category]].append((os.path.join(self.root, category, 'points', uuid+'.pts'),
                                        os.path.join(self.root, category, 'points_label', uuid+'.seg')))

        self.datapath = []
        for item in self.cat:
            for fn in self.meta[item]:
                self.datapath.append((item, fn[0], fn[1]))

        self.classes = dict(zip(sorted(self.cat), range(len(self.cat))))
        print(self.classes)
        with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../dataset/misc/num_seg_classes.txt'), 'r') as f:
            for line in f:
                ls = line.strip().split()
                self.seg_classes[ls[0]] = int(ls[1])
        self.num_seg_classes = self.seg_classes[list(self.cat.keys())[0]]
        print(self.seg_classes, self.num_seg_classes)

    def __getitem__(self, index):
        fn = self.datapath[index]
        cls = self.classes[self.datapath[index][0]]

        point_set = np.loadtxt(fn[1]).astype(np.float32)
        seg = np.loadtxt(fn[2]).astype(np.int64)
        print(point_set.shape, seg.shape)

        choice = np.random.choice(len(seg), self.npoints, replace=True)
        print(point_set.shape)
        #resample
        point_set = point_set[choice, :]
        print(" ######",point_set[0])
        point_set = point_set - np.expand_dims(np.mean(point_set, axis = 0), 0) # center
        dist = np.max(np.sqrt(np.sum(point_set ** 2, axis = 1)),0)
        point_set = point_set / dist #scale
        print(" $$$$$$$$$",point_set[0],dist)

        if self.data_augmentation:
            theta = np.random.uniform(0,np.pi*2)
            rotation_matrix = np.array([[np.cos(theta), -np.sin(theta)],[np.sin(theta), np.cos(theta)]])
            point_set[:,[0,2]] = point_set[:,[0,2]].dot(rotation_matrix) # random rotation
            point_set += np.random.normal(0, 0.02, size=point_set.shape) # random jitter

        seg = seg[choice]
        point_set = torch.from_numpy(point_set)
        seg = torch.from_numpy(seg)
        cls = torch.from_numpy(np.array([cls]).astype(np.int64))

        if self.classification:
            return point_set, cls
        else:
            return point_set, seg

    def __len__(self):
        return len(self.datapath)

模型设计

C++版：

比python多register_module(）model.h

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
//
// Created by sindre on 2022/1/10.
//

#ifndef LIBTORCH_TEMPLATE_MODEL_H
#define LIBTORCH_TEMPLATE_MODEL_H
#include <torch/torch.h>
using namespace torch::nn;
struct Result
{
    torch::Tensor x;
    torch::Tensor trans;
    torch::Tensor trans_feat;

};
class STN3dImpl :public torch::nn::Module {
private:
    Conv1d conv1= nullptr;
    Conv1d conv2= nullptr;
    Conv1d conv3= nullptr;
    Linear fc1= nullptr;
    Linear fc2= nullptr;
    Linear fc3= nullptr;
    ReLU relu= nullptr;

    BatchNorm1d bn1= nullptr;
    BatchNorm1d bn2= nullptr;
    BatchNorm1d bn3= nullptr;
    BatchNorm1d bn4= nullptr;
    BatchNorm1d bn5= nullptr;
public:
    STN3dImpl();
    torch::Tensor forward(torch::Tensor x);
};TORCH_MODULE(STN3d);

class STNkdImpl :public torch::nn::Module {
private:
    Conv1d conv1= nullptr;
    Conv1d conv2= nullptr;
    Conv1d conv3= nullptr;
    Linear fc1= nullptr;
    Linear fc2= nullptr;
    Linear fc3= nullptr;
    ReLU relu= nullptr;

    BatchNorm1d bn1= nullptr;
    BatchNorm1d bn2= nullptr;
    BatchNorm1d bn3= nullptr;
    BatchNorm1d bn4= nullptr;
    BatchNorm1d bn5= nullptr;

    int num=64;
public:
    STNkdImpl(int num);
    torch::Tensor forward(torch::Tensor x);
};TORCH_MODULE(STNkd);


class PointNetFeatImpl :public torch::nn::Module {
private:
    Conv1d conv1= nullptr;
    Conv1d conv2= nullptr;
    Conv1d conv3= nullptr;

    BatchNorm1d bn1= nullptr;
    BatchNorm1d bn2= nullptr;
    BatchNorm1d bn3= nullptr;

    STN3d stn= nullptr;
    STNkd fstn= nullptr;
    Result ret;
public:
    PointNetFeatImpl();
    Result forward(torch::Tensor x);
};TORCH_MODULE(PointNetFeat);


class PointNetSeg:public torch::nn::Module {
private:
    PointNetFeat feat= nullptr;
    Conv1d conv1= nullptr;
    Conv1d conv2= nullptr;
    Conv1d conv3= nullptr;
    Conv1d conv4= nullptr;

    BatchNorm1d bn1= nullptr;
    BatchNorm1d bn2= nullptr;
    BatchNorm1d bn3= nullptr;


    Result ret;
    Result ret_feat;

    int num=2;
public:
    PointNetSeg(int num=2);
    Result forward(torch::Tensor x);
};

#endif//LIBTORCH_TEMPLATE_MODEL_H

model.cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
//
// Created by sindre on 2022/1/10.
//

#include "../include/model.h"

STN3dImpl::STN3dImpl() {
    conv1 = Conv1d(torch::nn::Conv1dOptions(3,64,1));
    conv2 = Conv1d(torch::nn::Conv1dOptions(64,128,1));
    conv3 = Conv1d(torch::nn::Conv1dOptions(128,1024,1));
    fc1 = Linear(torch::nn::LinearOptions(1024,512));
    fc2 = Linear(torch::nn::LinearOptions(512,256));
    fc3= Linear(torch::nn::LinearOptions(256,9));
    relu = ReLU();



    bn1 = BatchNorm1d(torch::nn::BatchNorm1dOptions(64));
    bn2 = BatchNorm1d(torch::nn::BatchNorm1dOptions(128));
    bn3 = BatchNorm1d(torch::nn::BatchNorm1dOptions(1024));
    bn4 = BatchNorm1d(torch::nn::BatchNorm1dOptions(512));
    bn5 = BatchNorm1d(torch::nn::BatchNorm1dOptions(256));

    register_module("conv1",  conv1);
    register_module("conv2",  conv2);
    register_module("conv3",  conv3);
    register_module("fc1",  fc1);
    register_module("fc2",  fc2);
    register_module("fc3",  fc3);
    register_module("ReLU",  relu);

    register_module("bn1", bn1);
    register_module("bn2", bn2);
    register_module("bn3", bn3);
    register_module("bn4", bn4);
    register_module("bn5", bn5);

}
torch::Tensor STN3dImpl::forward(torch::Tensor x) {
    int batchsize =x.size(0);
    x=torch::relu(bn1->forward(conv1->forward(x)));
    //std::cout <<"out"<<x[0]<< std::endl;
    x=torch::relu(bn2->forward(conv2->forward(x)));
    x=torch::relu(bn3->forward(conv3->forward(x)));

    x=std::get<0>(torch::max(x,2,true));
    x=x.view({-1,1024});

    x=torch::relu(bn4->forward(fc1->forward(x)));
    x=torch::relu(bn5->forward(fc2->forward(x)));
    x=fc3->forward(x);

    torch::Tensor iden = torch::tensor({1,0,0,0,1,0,0,0,1}).view({1,9}).repeat({batchsize,1});
    if(x.is_cuda()){
        iden = iden.cuda();
    }
    x = x+iden;
    x = x.view({-1,3,3});
    return x;
}


STNkdImpl::STNkdImpl(int num) {
    conv1 = Conv1d(Conv1dOptions(num,64,1));
    conv2 = Conv1d(Conv1dOptions(64,128,1));
    conv3 = Conv1d(Conv1dOptions(128,1024,1));
    fc1 = Linear(1024,512);
    fc2 = Linear(512,256);
    fc3= Linear(256,num*num);
    relu = ReLU();



    bn1 = BatchNorm1d(64);
    bn2 = BatchNorm1d(128);
    bn3 = BatchNorm1d(1024);
    bn4 = BatchNorm1d(512);
    bn5 = BatchNorm1d(256);

    STNkdImpl::num = num;

    register_module("conv1",  conv1);
    register_module("conv2",  conv2);
    register_module("conv3",  conv3);
    register_module("fc1",  fc1);
    register_module("fc2",  fc2);
    register_module("fc3",  fc3);
    register_module("ReLU",  relu);

    register_module("bn1", bn1);
    register_module("bn2", bn2);
    register_module("bn3", bn3);
    register_module("bn4", bn4);
    register_module("bn5", bn5);

}
torch::Tensor STNkdImpl::forward(torch::Tensor x) {

    int batchsize =x.size(0);
    x=torch::relu(bn1->forward(conv1->forward(x)));
    x=torch::relu(bn2->forward(conv2->forward(x)));
    x=torch::relu(bn3->forward(conv3->forward(x)));
    x=torch::amax(x,2,true);
    x=x.view({-1,1024});

    x=torch::relu(bn4->forward(fc1->forward(x)));
    x=torch::relu(bn5->forward(fc2->forward(x)));
    x=fc3->forward(x);

    torch::Tensor iden = torch::eye(STNkdImpl::num).view({1,STNkdImpl::num*STNkdImpl::num}).repeat({batchsize,1});
    if(x.is_cuda()){
        iden = iden.cuda();
    }

    x = x+iden;
    x = x.view({-1,STNkdImpl::num,STNkdImpl::num});
    return x;
}

PointNetFeatImpl::PointNetFeatImpl() {
    stn = STN3d();

    conv1 = Conv1d(Conv1dOptions(3,64,1));
    conv2 = Conv1d(Conv1dOptions(64,128,1));
    conv3 = Conv1d(Conv1dOptions(128,1024,1));
    bn1 = BatchNorm1d(64);
    bn2 = BatchNorm1d(128);
    bn3 = BatchNorm1d(1024);

    fstn = STNkd(64);

    register_module("stn",  stn );
    register_module("conv1",  conv1);
    register_module("conv2",  conv2);
    register_module("conv3",  conv3);
    register_module("bn1",  bn1);
    register_module("bn2",  bn2);
    register_module("bn3",  bn3);
    register_module("fstn",  fstn);


}


Result PointNetFeatImpl::forward(torch::Tensor x) {
    int64_t n_pts = x.size(2);
    torch::Tensor trans =stn->forward(x);

    x = x.transpose(2,1);
    x = torch::bmm(x,trans);
    x = x.transpose(2,1);
    x = torch::relu(bn1->forward(conv1->forward(x)));

    torch::Tensor trans_feat = fstn->forward(x);

    x = x.transpose(2,1);
    x = torch::bmm(x,trans_feat);
    x = x.transpose(2,1);


    torch::Tensor pointfeat = x ;
    x = torch::relu(bn2->forward(conv2->forward(x)));
    x = bn3->forward(conv3->forward(x));
    x = torch::amax(x,2,true);
    x = x.view({-1,1024});
    //只做分割
    x = x.view({-1,1024,1}).repeat({1,1,n_pts});



    ret.trans_feat=trans_feat;
    ret.trans=trans;
    ret.x =torch::cat({x,pointfeat},1);
    return ret;

}


PointNetSeg::PointNetSeg(int num) {
    PointNetSeg::num=num;
    feat =PointNetFeat();
    conv1 = Conv1d(Conv1dOptions(1088,512,1));
    conv2 = Conv1d(Conv1dOptions(512,256,1));
    conv3 = Conv1d(Conv1dOptions(256,128,1));
    conv4 = Conv1d(Conv1dOptions(128,num,1));
    bn1 = BatchNorm1d(512);
    bn2 = BatchNorm1d(256);
    bn3 = BatchNorm1d(128);

    register_module("feat",   feat );
    register_module("conv1",  conv1);
    register_module("conv2",  conv2);
    register_module("conv3",  conv3);
    register_module("conv4",  conv4);
    register_module("bn1",  bn1);
    register_module("bn2",  bn2);
    register_module("bn3",  bn3);

}
Result PointNetSeg::forward(torch::Tensor x) {
    int64_t batchsize =x.size(0);
    int64_t n_pts = x.size(2);
    int64_t num_ =PointNetSeg::num;


    ret_feat = feat->forward(x);
    x=ret_feat.x;



    x=torch::relu(bn1->forward(conv1->forward(x)));
    x=torch::relu(bn2->forward(conv2->forward(x)));
    x=torch::relu(bn3->forward(conv3->forward(x)));
    x=conv4->forward(x);
    x=x.transpose(2,1).contiguous();
    x=torch::log_softmax(x.view({-1,num_}),-1);
    x=x.view({batchsize,n_pts,num_});



    ret.trans_feat=ret_feat.trans_feat;
    ret.trans=ret_feat.trans;
    ret.x =x;
    return ret;
}

python版：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
from torch.autograd import Variable
import numpy as np
import torch.nn.functional as F


class STN3d(nn.Module):
    def __init__(self):
        super(STN3d, self).__init__()
        self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 9)
        self.relu = nn.ReLU()

        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(256)


    def forward(self, x):
        batchsize = x.size()[0]
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)

        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(self.fc2(x)))
        x = self.fc3(x)

        iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1)
        if x.is_cuda:
            iden = iden.cuda()
        x = x + iden
        x = x.view(-1, 3, 3)
        return x


class STNkd(nn.Module):
    def __init__(self, k=64):
        super(STNkd, self).__init__()
        self.conv1 = torch.nn.Conv1d(k, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k*k)
        self.relu = nn.ReLU()

        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(256)

        self.k = k

    def forward(self, x):
        batchsize = x.size()[0]
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)

        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(self.fc2(x)))
        x = self.fc3(x)

        iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1,self.k*self.k).repeat(batchsize,1)
        if x.is_cuda:
            iden = iden.cuda()
        x = x + iden
        x = x.view(-1, self.k, self.k)
        return x

class PointNetfeat(nn.Module):
    def __init__(self, global_feat = True, feature_transform = False):
        super(PointNetfeat, self).__init__()
        self.stn = STN3d()
        self.conv1 = torch.nn.Conv1d(3, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.global_feat = global_feat
        self.feature_transform = feature_transform
        if self.feature_transform:
            self.fstn = STNkd(k=64)

    def forward(self, x):
        n_pts = x.size()[2]
        trans = self.stn(x)
        x = x.transpose(2, 1)
        x = torch.bmm(x, trans)
        x = x.transpose(2, 1)
        x = F.relu(self.bn1(self.conv1(x)))

        if self.feature_transform:
            trans_feat = self.fstn(x)
            x = x.transpose(2,1)
            x = torch.bmm(x, trans_feat)
            x = x.transpose(2,1)
        else:
            trans_feat = None

        pointfeat = x
        x = F.relu(self.bn2(self.conv2(x)))
        x = self.bn3(self.conv3(x))
        x = torch.max(x, 2, keepdim=True)[0]
        x = x.view(-1, 1024)
        if self.global_feat:
            return x, trans, trans_feat
        else:
            x = x.view(-1, 1024, 1).repeat(1, 1, n_pts)
            return torch.cat([x, pointfeat], 1), trans, trans_feat

class PointNetCls(nn.Module):
    def __init__(self, k=2, feature_transform=False):
        super(PointNetCls, self).__init__()
        self.feature_transform = feature_transform
        self.feat = PointNetfeat(global_feat=True, feature_transform=feature_transform)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k)
        self.dropout = nn.Dropout(p=0.3)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.relu = nn.ReLU()

    def forward(self, x):
        x, trans, trans_feat = self.feat(x)
        x = F.relu(self.bn1(self.fc1(x)))
        x = F.relu(self.bn2(self.dropout(self.fc2(x))))
        x = self.fc3(x)
        return F.log_softmax(x, dim=1), trans, trans_feat


class PointNetDenseCls(nn.Module):
    def __init__(self, k = 2, feature_transform=False):
        super(PointNetDenseCls, self).__init__()
        self.k = k
        self.feature_transform=feature_transform
        self.feat = PointNetfeat(global_feat=False, feature_transform=feature_transform)
        self.conv1 = torch.nn.Conv1d(1088, 512, 1)
        self.conv2 = torch.nn.Conv1d(512, 256, 1)
        self.conv3 = torch.nn.Conv1d(256, 128, 1)
        self.conv4 = torch.nn.Conv1d(128, self.k, 1)
        self.bn1 = nn.BatchNorm1d(512)
        self.bn2 = nn.BatchNorm1d(256)
        self.bn3 = nn.BatchNorm1d(128)

    def forward(self, x):
        #print("xxx：",x.size())
        batchsize = x.size()[0]
        n_pts = x.size()[2]
        x, trans, trans_feat = self.feat(x)
        #print("xxx：",x.size())
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.conv4(x)
        x = x.transpose(2,1).contiguous()
        x = F.log_softmax(x.view(-1,self.k), dim=-1)


        #print("xxx：",x.size())
        x = x.view(batchsize, n_pts, self.k)

        return x, trans, trans_feat

def feature_transform_regularizer(trans):
    d = trans.size()[1]
    batchsize = trans.size()[0]
    I = torch.eye(d)[None, :, :]
    if trans.is_cuda:
        I = I.cuda()
    loss = torch.mean(torch.norm(torch.bmm(trans, trans.transpose(2,1)) - I, dim=(1,2)))
    return loss

if __name__ == '__main__':
    sim_data = Variable(torch.rand(32,3,2500))
    trans = STN3d()
    out = trans(sim_data)
    print('stn', out.size())
    print('loss', feature_transform_regularizer(out))

    sim_data_64d = Variable(torch.rand(32, 64, 2500))
    trans = STNkd(k=64)
    out = trans(sim_data_64d)
    print('stn64d', out.size())
    print('loss', feature_transform_regularizer(out))

    pointfeat = PointNetfeat(global_feat=True)
    out, _, _ = pointfeat(sim_data)
    print('global feat', out.size())

    pointfeat = PointNetfeat(global_feat=False)
    out, _, _ = pointfeat(sim_data)
    print('point feat', out.size())

    cls = PointNetCls(k = 5)
    out, _, _ = cls(sim_data)
    print('class', out.size())

    seg = PointNetDenseCls(k = 3)
    out, _, _ = seg(sim_data)
    print('seg', out.size())
    print(out)

训练步骤

C++版：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
//
// Created by sindre on 2022/1/10.
//

#ifndef LIBTORCH_TEMPLATE_TRAIN_H
#define LIBTORCH_TEMPLATE_TRAIN_H

#include "../include/dataset.h"
#include "../include/model.h"
#include "opencv2/opencv.hpp"
#include <torch/torch.h>
#include <fmt/core.h>
#include <string>

#include "torch/script.h"

using namespace torch::indexing;
using namespace fmt;
using namespace std;

class train {
private:
    size_t batchsize = 60;
    size_t workers = 8;
    size_t nepoch = 250 ;
    int npoints = 2500 ;
    int cls = 2;//数据标签只有两类
    std::string model_path ;
    std::string dataset_path;


public:
    train(size_t _batchsize, size_t _workers,size_t _nepoch,std::string _model,std::string _dataset);

};


#endif//LIBTORCH_TEMPLATE_TRAIN_H

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
//
// Created by sindre on 2022/1/10.
//

#include "../include/train.h"
namespace F = torch::nn::functional;
train::train(size_t _batchsize, size_t _workers, size_t _nepoch, std::string _model, std::string _dataset_path) {
    batchsize = _batchsize;
    workers = _workers;
    model_path = _model;
    dataset_path = _dataset_path;
    auto device = torch::Device(torch::kCPU);
    if(torch::cuda::is_available()){device =torch::Device(torch::kCUDA,0); }

    //固定随机种子
    auto manualSeed = rand()%1000;
    print("seed : {} \n",manualSeed);
    torch::manual_seed(manualSeed);

    //加载数据
    auto train_set = dataset(dataset_path,npoints).map(torch::data::transforms::Stack<>());
    auto train_load = torch::data::make_data_loader<torch::data::samplers::SequentialSampler>(std::move(train_set),batchsize);
    print("dataset load done \n");

    //加载模型
    auto Seg_model = PointNetSeg(cls);
    Seg_model.to(device);
    print("Seg_model load done \n");
    print("{}\n",Seg_model);


    //优化器
    auto optimizer = torch::optim::Adam(Seg_model.parameters(),0.001);


    //训练
    print("************start training*************\n");
    for(int i=0; i<nepoch; i++) {
        for (auto &batch: *train_load) {
            torch::Tensor data = batch.data;
            //torch::Tensor data = torch::rand({12,3,2500});
            torch::Tensor target = batch.target;

            data = data.to(device);
            target = target.to(device);
            optimizer.zero_grad();
            Seg_model.train();
//            print("************complete initialization*******data:{}****target: {}**\n",data.sizes(),target.sizes());

            auto seg_out = Seg_model.forward(data);
//            print("modle out: {}\n",seg_out.x.sizes());
            torch::Tensor pred = seg_out.x.view({-1, cls});
            target = target.view({-1, 1}).index({"...",0})-1;
//            print(target[0]);
            target = target.to(torch::kLong);
//            print("***p:{}--->label:{}  \n type: {} {}*******compute loss*************\n",pred.sizes(),target.sizes(),pred.dtype(),target.dtype());


            //pred = torch::rand({30000,2});
            //target = torch::rand({30000}).to(torch::kLong);
            //pred.print();
            //target.print();
            torch::Tensor loss = torch::nll_loss(pred, target);
            //loss.print();

            //loss += seg_out.trans_feat * 0.001;
            //print("**********complete loss*********************");
            loss.backward();
            optimizer.step();
            print("[{}/{}] train loss:{} \n", nepoch, i, loss.item());
        }
        torch::save(Seg_model.parameters(), "seg_model.pt");
    }



}

python版：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
from __future__ import print_function




import argparse
import os
import random

import numpy as np
import torch.nn.functional as F
import torch.utils.data
from tqdm import tqdm

# from pointnet.dataset import ShapeNetDataset
from pointnet.model import PointNetDenseCls, feature_transform_regularizer
from user_dataset import  User_Dataset

parser = argparse.ArgumentParser()
parser.add_argument(
    '--batchSize', type=int, default=16, help='input batch size')
parser.add_argument(
    '--workers', type=int, help='number of data loading workers', default=8)
parser.add_argument(
    '--nepoch', type=int, default=250, help='number of epochs to train for')
parser.add_argument('--model', type=str, default='', help='model path')
parser.add_argument('--dataset', type=str, default='../dataset/shapenetcore', help="dataset path")
parser.add_argument('--class_choice', type=str, default='Chair', help="class_choice")
parser.add_argument('--feature_transform', action='store_true', help="use feature transform")

opt = parser.parse_args()
print(opt)

opt.manualSeed = random.randint(1, 10000)  # fix seed
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)

dataset = User_Dataset()
dataloader = torch.utils.data.DataLoader(
    dataset,
    batch_size=opt.batchSize,
    shuffle=True,
    num_workers=int(opt.workers))

test_dataset = User_Dataset(split='test')
testdataloader = torch.utils.data.DataLoader(
    test_dataset,
    batch_size=opt.batchSize,
    shuffle=True,
    num_workers=int(opt.workers))


blue = lambda x: '\033[94m' + x + '\033[0m'
num_classes=2 #数据只需要分割成两类
classifier = PointNetDenseCls(k=num_classes, feature_transform=opt.feature_transform)
print(classifier)
if opt.model != '':
    classifier.load_state_dict(torch.load(opt.model))

optimizer = torch.optim.Adam(classifier.parameters(), lr=0.001, betas=(0.9, 0.999))
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)
classifier.cuda()

num_batch = len(dataset) / opt.batchSize
score_record=0
for epoch in range(opt.nepoch):
    scheduler.step()

    for i, data in enumerate(dataloader):
        points, target = data
        points = points.transpose(2, 1)
        points, target = points.cuda(), target.cuda()
        optimizer.zero_grad()
        classifier = classifier.train()
        pred, trans, trans_feat = classifier(points)
        pred = pred.view(-1, num_classes)
        target = target.view(-1, 1)[:, 0] - 1
        #print(pred.size(), target.size())
        loss = F.nll_loss(pred, target)
        if opt.feature_transform:
            loss += feature_transform_regularizer(trans_feat) * 0.001
        loss.backward()
        optimizer.step()
        pred_choice = pred.data.max(1)[1]
        correct = pred_choice.eq(target.data).cpu().sum()
        print('[%d: %d/%d] train loss: %f accuracy: %f' % (epoch, i, num_batch, loss.item(), correct.item()/float(opt.batchSize * 2500)))

        if i % 10 == 0:
            j, data = next(enumerate(testdataloader, 0))
            points, target = data
            points = points.transpose(2, 1)
            points, target = points.cuda(), target.cuda()
            classifier = classifier.eval()
            pred, _, _ = classifier(points)
            pred = pred.view(-1, num_classes)
            target = target.view(-1, 1)[:, 0] - 1



            loss = F.nll_loss(pred, target)
            pred_choice = pred.data.max(1)[1]
            correct = pred_choice.eq(target.data).cpu().sum()
            print('[%d: %d/%d] %s loss: %f accuracy: %f' % (epoch, i, num_batch, blue('test'), loss.item(), correct.item()/float(opt.batchSize * 2500)))
            score_record=correct.item()
    #if correct.item()>score_record:
    torch.save(classifier.state_dict(), 'seg_model.pth')

## benchmark mIOU
shape_ious = []
for i,data in tqdm(enumerate(testdataloader, 0)):
    points, target = data
    points = points.transpose(2, 1)
    points, target = points.cuda(), target.cuda()
    classifier = classifier.eval()
    pred, _, _ = classifier(points)
    pred_choice = pred.data.max(2)[1]

    pred_np = pred_choice.cpu().data.numpy()
    target_np = target.cpu().data.numpy() - 1

    for shape_idx in range(target_np.shape[0]):
        parts = range(num_classes)#np.unique(target_np[shape_idx])
        part_ious = []
        for part in parts:
            I = np.sum(np.logical_and(pred_np[shape_idx] == part, target_np[shape_idx] == part))
            U = np.sum(np.logical_or(pred_np[shape_idx] == part, target_np[shape_idx] == part))
            if U == 0:
                iou = 1 #If the union of groundtruth and prediction points is empty, then count part IoU as 1
            else:
                iou = I / float(U)
            part_ious.append(iou)
        shape_ious.append(np.mean(part_ious))

print("mIOU for class {}: {}".format(opt.class_choice, np.mean(shape_ious)))


相关资料

链接：https://up3dai.yuque.com/docs/share/9cef4df7-f676-4bb5-acd9-95a97448b403?# 《PointNet:用于点云的分类和分割深度学习》https://up3dai.yuque.com/docs/share/13e72a10-a5da-4a0b-8d7e-1e9c8aecd6e9?# 《PointNet:pytorch版代码解析》

欢迎关注我的其它发布渠道