点云深度学习系列博客(四)：PointNet代码精讲

目录

1. 代码解析

1.1 初始化

1.2 数据载入

1.3 模型载入

1.4 训练代码

2. 实验结果

Reference

---

最近开始上手点云深度学习项目，相比之前纸上谈兵的阶段，此时我将把更多的精力放在代码学习和复现上。在新的学习阶段，就不能是看看论文，蜻蜓点水的配下别人的代码这么简单了。我将逐句分析代码功能，结合实际应用，来深入理解点云深度学习的项目该如何落地。作为点云深度学习的代表作，PointNet [1] 的经典程度不言而喻。我们就以PointNet的模板，来展开相关代码的实现，并完全复现PointNet的基本功能。对于那些计划零基础入坑点云深度学习的同学，不妨看看。

---

1. 代码解析

我们的解析基于项目：https://github.com/yanx27/Pointnet_Pointnet2_pytorch

使用的平台为Pycharm2021+python3.8+pytorch12.1+cu116。

我们以train_classification.py项目为例，来介绍代码的实现细节。这里使用的数据库为ModelNet。

1.1 初始化

首先我们看看在正式训练之前，程序都做了什么事情。按照项目文档，参数化调用：

python train_classification.py --model pointnet2_cls_ssg --log_dir pointnet2_cls_ssg

按照函数parse_args()进行解析，具体为：

--use_cpu cpu模式选择

--gpu gpu模式选择

--batch_size 数据块尺度

--model 指定训练模型

--num_category 分类数目

--epoch 数据扫描次数

--learning_rate 学习率

--num_point 点云采样点数

--optimizer 优化器，默认选择

--log_dir 实验根目录

--decay_rate 衰减率

--use_normals 是否使用法线

--process_data 是否离线处理数据

--use_uniform_sample 是否使用均匀采样

标红的项为对应之前参数化调用的两项。接下来，我们看main函数的内容。首先，我们介绍获取参数调用后的一些初始设置，包括路径设置以及日志文件设置。

#存储可用的gpu
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
 
timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M'))
 
#设置日志文件路径
exp_dir = Path('./log/')
exp_dir.mkdir(exist_ok=True)
exp_dir = exp_dir.joinpath('classification')
exp_dir.mkdir(exist_ok=True)
 
#建立log路径，如果没有给定路径名，建立以时间为名的路径
if args.log_dir is None:
    exp_dir = exp_dir.joinpath(timestr)
else:
    exp_dir = exp_dir.joinpath(args.log_dir)
exp_dir.mkdir(exist_ok=True)
 
checkpoints_dir = exp_dir.joinpath('checkpoints/')
checkpoints_dir.mkdir(exist_ok=True)
log_dir = exp_dir.joinpath('logs/')
log_dir.mkdir(exist_ok=True)
 
#这一部分用于存储log信息，将时间以及相关参数存在（log_dir/args.model).txt中
logger = logging.getLogger("Model")
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model))
file_handler.setLevel(logging.INFO)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
log_string('PARAMETER ...')
log_string(args)

1.2 数据载入

train_dataset = ModelNetDataLoader(root=data_path, args=args, split='train', 
                process_data=args.process_data)
test_dataset = ModelNetDataLoader(root=data_path, args=args, split='test', 
               process_data=args.process_data)
trainDataLoader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, 
                  shuffle=True, num_workers=10, drop_last=True)
testDataLoader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, 
                 shuffle=False, num_workers=10)

这里展开解释一下ModelNetDataLoader：

    def __init__(self, root, args, split='train', process_data=False):
        
        #首先初始化参数，包括根目录root，参数列表args，训练集或测试集指定，默认是训练集，以及 
        是否需要离线存储数据process_data
 
        self.root = root        
        self.process_data = process_data
 
        #从参数列表中提取参数，包括点数，是否均匀化采样，是否使用法线，分类数
        self.npoints = args.num_point
        self.uniform = args.use_uniform_sample
        self.use_normals = args.use_normals
        self.num_category = args.num_category
        
        #这里判断下，如果是10类，就读取modelnet10_shape_names.txt，否则读取 
        modelnet40_shape_names.txt
        if self.num_category == 10:
            self.catfile = os.path.join(self.root, 'modelnet10_shape_names.txt')
        else:
            self.catfile = os.path.join(self.root, 'modelnet40_shape_names.txt')
 
        #这里的self.cat类型是list，存储的是类别
        #这里的self.classes 类型是dict，存储的是类别和对应的编号
        self.cat = [line.rstrip() for line in open(self.catfile)]
        self.classes = dict(zip(self.cat, range(len(self.cat))))
 
        #这里的shape_ids类型是dict，存储的是train的list和test的list
        shape_ids = {}
        if self.num_category == 10:
            shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 
            'modelnet10_train.txt'))]
            shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 
            'modelnet10_test.txt'))]
        else:
            shape_ids['train'] = [line.rstrip() for line in open(os.path.join(self.root, 
            'modelnet40_train.txt'))]
            shape_ids['test'] = [line.rstrip() for line in open(os.path.join(self.root, 
            'modelnet40_test.txt'))]
 
        assert (split == 'train' or split == 'test')
 
        #从shape_ids这个dict中，按照split的指定提取shape名
        shape_names = ['_'.join(x.split('_')[0:-1]) for x in shape_ids[split]]
 
        #按照shape名设置父路径，并且把对应的model存储在对应的路径里
        self.datapath = [(shape_names[i], os.path.join(self.root, shape_names[i], 
        shape_ids[split][i]) + '.txt') for i
                         in range(len(shape_ids[split]))]
        print('The size of %s data is %d' % (split, len(self.datapath)))
 
        #这两句不知道什么意思，好像是设置了一个关于是否均匀采样的路径。但是我运行后，对应路径 
        #没有发现.dat文件
        if self.uniform:
            self.save_path = os.path.join(root, 'modelnet%d_%s_%dpts_fps.dat' % 
            (self.num_category, split, self.npoints))
        else:
            self.save_path = os.path.join(root, 'modelnet%d_%s_%dpts.dat' % 
            (self.num_category, split, self.npoints))
 
        #下面相当长的一段是用来做离线数据处理的，主要用来统一点数，如果这里不做处理，那么就是 
        #在训练的时候做处理。
        if self.process_data:
            if not os.path.exists(self.save_path):
                print('Processing data %s (only running in the first time)...' % 
                self.save_path)
                self.list_of_points = [None] * len(self.datapath)
                self.list_of_labels = [None] * len(self.datapath)
 
                for index in tqdm(range(len(self.datapath)), total=len(self.datapath)):
                    fn = self.datapath[index]
                    cls = self.classes[self.datapath[index][0]]
                    cls = np.array([cls]).astype(np.int32)
                    point_set = np.loadtxt(fn[1], delimiter=',').astype(np.float32)
 
                    if self.uniform:
                        point_set = farthest_point_sample(point_set, self.npoints)
                    else:
                        point_set = point_set[0:self.npoints, :]
 
                    self.list_of_points[index] = point_set
                    self.list_of_labels[index] = cls
 
                with open(self.save_path, 'wb') as f:
                    pickle.dump([self.list_of_points, self.list_of_labels], f)
            else:
                print('Load processed data from %s...' % self.save_path)
                with open(self.save_path, 'rb') as f:
                    self.list_of_points, self.list_of_labels = pickle.load(f)

我们知道，在进行网络训练的时候，为了能够在有限的存储结构基础上，实现对大规模数据集的训练，需要把数据打包成batch，然后再每一个batch上进行运算，并在最后组合不同的batch。对于Pontnet项目，给出的默认batch_size=24。那么，我们在训练的时候，需要把前一阶段载入的数据，按照batch_size进行整理，以方便网络使用，这里，项目使用torch.utils.data.DataLoader来完成相关功能，代码如下：

trainDataLoader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, 
                  shuffle=True, num_workers=10, drop_last=True)
testDataLoader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size, 
                 shuffle=False, num_workers=10)

shuffle（洗牌）表示是否在每一次训练时，重组数据；number_workers表示工作者数量，默认是0。使用多少个子进程来导入数据 drop_last表示是否丢弃最后一个不完整batch的数据。

1.3 模型载入

#载入模型
num_class = args.num_category
#动态导入对象，这里默认的是pointnet_cls
model = importlib.import_module(args.model)
 
#没啥用的shutil，建议删掉
shutil.copy('./models/%s.py' % args.model, str(exp_dir))
shutil.copy('models/pointnet2_utils.py', str(exp_dir))
shutil.copy('./train_classification.py', str(exp_dir))
 
#这里调用的get_model，即pointnet的网络结构
classifier = model.get_model(num_class, normal_channel=args.use_normals)
criterion = model.get_loss()
 
#这里判断了一下激活函数是否是ReLu，如果不是，设置成ReLu
classifier.apply(inplace_relu)

不知道作者使用shutil.copy的原因，要把三个文件在log复制一遍。我在项目中把这个部分删掉了，试了一下，没有任何影响。importlib.import_module为动态导入对象。model.get_model即为提取pointnet的backbone。因为这里涉及到了model最核心的代码，所以要展开，如下：

    def __init__(self, k=40, normal_channel=True):
        super(get_model, self).__init__()
        if normal_channel:
            channel = 6
        else:
            channel = 3
        self.feat = PointNetEncoder(global_feat=True, feature_transform=True, 
                    channel=channel)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k)
        self.dropout = nn.Dropout(p=0.4)
        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)
        x = F.log_softmax(x, dim=1)
        return x, trans_feat

这20多行代码就是整个pointnet的backbone。可以看到结构是非常清晰的。这里项目的作者在输出最后分类的结果之前的MLP加了一个dropout，并把参数设为0.4。这是一个trick，用来减少过拟合。PointNetEncoder已经封装了自maxpooling以前的全部代码。因此，只需要加三个mlp全连接层，以补全网络的后段实现就可以。在pointnet_utils.py中存储着PointNetEncoder的结构：

class PointNetEncoder(nn.Module):
    def __init__(self, global_feat=True, feature_transform=False, channel=3):
        super(PointNetEncoder, self).__init__()
        #PointNet的第一个T-Net
        self.stn = STN3d(channel)
        
        #使用一维卷积实现MLP，从64升维到1024
        self.conv1 = torch.nn.Conv1d(channel, 64, 1)
        self.conv2 = torch.nn.Conv1d(64, 128, 1)
        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
 
        #参数归一化，配合MLP
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        
        #这个用来判断是否复制全局特征（maxpooling后的1024维全局特征，分类问题就不用复制，分割 
        问题就需要复制，pointnet_sem_seg的设置就是需要复制的）
        self.global_feat = global_feat
        
        #按照参数设置，以判断是否需要特征对齐
        self.feature_transform = feature_transform
        if self.feature_transform:
            self.fstn = STNkd(k=64)
 
    def forward(self, x):
        B, D, N = x.size()
        trans = self.stn(x)
        x = x.transpose(2, 1)
        if D > 3:
            feature = x[:, :, 3:]
            x = x[:, :, :3]
 
        #从stn中获取变换矩阵trans，使用bmm批量变换x
        x = torch.bmm(x, trans)
        if D > 3:
            x = torch.cat([x, feature], dim=2)
        
        #转置
        x = x.transpose(2, 1)
 
        #这里是第一个T-Net之后的MLP
        x = F.relu(self.bn1(self.conv1(x)))
 
        #这里判断是否使用第二个T-Net
        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
 
        #这里是第二个T-Net之后的MLP
        x = F.relu(self.bn2(self.conv2(x)))
 
        #这里是第三个MLP，此时每个点到了1024维
        x = self.bn3(self.conv3(x))
 
        #maxpooing
        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)
            return torch.cat([x, pointfeat], 1), trans, trans_feat

1.4 训练代码

让我们回到主程序，在完成model的载入后，我们需要做一些基本的设置，就可以开始训练环节了，包括GPU模式确认，中间结果存储设置，优化器选择（Adam或SGD），设置调整学习率的机制，并存储在scheduler。然后设置存储分类结果的变量，以显示每一轮的分类精度：

#set device model
    if not args.use_cpu:
        classifier = classifier.cuda()
        criterion = criterion.cuda()
 
    try:
        checkpoint = torch.load(str(exp_dir) + '/checkpoints/best_model.pth')
        start_epoch = checkpoint['epoch']
        classifier.load_state_dict(checkpoint['model_state_dict'])
        log_string('Use pretrain model')
    except:
        log_string('No existing model, starting training from scratch...')
        start_epoch = 0
 
    if args.optimizer == 'Adam':
        optimizer = torch.optim.Adam(
            classifier.parameters(),
            lr=args.learning_rate,
            betas=(0.9, 0.999),
            eps=1e-08,
            weight_decay=args.decay_rate
        )
    else:
        optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9)
 
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.7)
    global_epoch = 0
    global_step = 0
    best_instance_acc = 0.0
    best_class_acc = 0.0

接下来，就是正式的训练代码：

    #开始扫描数据，这里的start_epoch记录的是之前训练的次数，首次训练为0
    for epoch in range(start_epoch, args.epoch):
        log_string('Epoch %d (%d/%s):' % (global_epoch + 1, epoch + 1, args.epoch))
        mean_correct = []
        #model开启训练模式
        classifier = classifier.train()
        """在scheduler的step_size表示scheduler.step()每调用step_size次，对应的学习率就会按 
           照策略调整一次。所以如果scheduler.step()是放在mini-batch里面，那么step_size指的 
           是经过这么多次迭代，学习率改变一次。"""
        scheduler.step()
        for batch_id, (points, target) in tqdm(enumerate(trainDataLoader, 0), 
            total=len(trainDataLoader), smoothing=0.9):
            #这里的points，target对应ModelNetDataLoader的_get_item方法，获取的是point_set和 
            #label[0]
             
            #梯度归0
            optimizer.zero_grad()
            
            #归一化点云尺度
            points = points.data.numpy()
            points = provider.random_point_dropout(points)
            points[:, :, 0:3] = provider.random_scale_point_cloud(points[:, :, 0:3])
            points[:, :, 0:3] = provider.shift_point_cloud(points[:, :, 0:3])
            points = torch.Tensor(points)
            points = points.transpose(2, 1)
 
            if not args.use_cpu:
                points, target = points.cuda(), target.cuda()
 
            #点云放入模型开始训练
            pred, trans_feat = classifier(points)
            #求loss
            loss = criterion(pred, target.long(), trans_feat)
            #预测分类结果
            pred_choice = pred.data.max(1)[1]
            #计算正确率
            correct = pred_choice.eq(target.long().data).cpu().sum()
            mean_correct.append(correct.item() / float(points.size()[0]))
            #梯度反传
            loss.backward()
            #更新学习率
            optimizer.step()
            global_step += 1
        
        #得到训练数据集的准确率
        train_instance_acc = np.mean(mean_correct)
        log_string('Train Instance Accuracy: %f' % train_instance_acc)
 
        with torch.no_grad():
 
            #使用模型来对测试数据进行测试。
            instance_acc, class_acc = test(classifier.eval(), testDataLoader, 
            num_class=num_class)
 
            #赋值best_instance_acc
            if (instance_acc >= best_instance_acc):
                best_instance_acc = instance_acc
                best_epoch = epoch + 1            
            if (class_acc >= best_class_acc):
                best_class_acc = class_acc
            log_string('Test Instance Accuracy: %f, Class Accuracy: %f' % (instance_acc, 
            class_acc))
            log_string('Best Instance Accuracy: %f, Class Accuracy: %f' % 
            (best_instance_acc, best_class_acc))
            
            #如果更新后，得到的精度高于之前的模型，存储，否则，不存储。
            if (instance_acc >= best_instance_acc):
                logger.info('Save model...')
                savepath = str(checkpoints_dir) + '/best_model.pth'
                log_string('Saving at %s' % savepath)
                state = {
                    'epoch': best_epoch,
                    'instance_acc': instance_acc,
                    'class_acc': class_acc,
                    'model_state_dict': classifier.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                }
                torch.save(state, savepath)
            global_epoch += 1
 
    logger.info('End of training...')

2. 实验结果

这里，我们设置epoch=100，并且直接贴出结果：

epoch 10：

Test Instance Accuracy: 0.762298, Class Accuracy: 0.693224
Best Instance Accuracy: 0.762298, Class Accuracy: 0.693224

epoch 50：

Test Instance Accuracy: 0.876214, Class Accuracy: 0.817101
Best Instance Accuracy: 0.878074, Class Accuracy: 0.822420

epoch 100：

Test Instance Accuracy: 0.887621, Class Accuracy: 0.833796
Best Instance Accuracy: 0.890129, Class Accuracy: 0.850198

可以看到，epoch 100后的分类准确率为0.89。

Reference

[1] Qi C, et al. Pointnet: Deep learning on point sets for 3d classification and segmentation[C]. Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 652-660.