使用Colab进行点云语义分割模型训练

前言

由于Colab免费版的内存过小，本文基于Colab Pro运行。
使用代码https://github.com/DLW3D/point_edg由https://github.com/yanx27/Pointnet_Pointnet2_pytorch修改而来。
可以直接使用Colab记事本打开。

运行步骤

Colab准备

加载硬盘，用于储存模型。

from google.colab import drive
drive.mount('/content/drive')

查看版本

!python --version
!cat /usr/local/cuda/version.txt
%tensorflow_version 1.x
import tensorflow as tf
print('Tensorflow version: {}'.format(tf.__version__))
import torch
print('Pytorch version: {}'.format(torch.__version__))
!nvidia-smi -L
print('GPU Identified at: {}'.format(tf.test.gpu_device_name()))

下载储存库

!git clone --depth=1 https://github.com/DLW3D/point_edg.git

下载S3DIS数据集

!mkdir -p /content/point_edg/data/s3dis

cd /content/point_edg/data/s3dis

!gdown --id 0BweDykwS9vIoaTRsbWpnWjFfUU0

解压数据集，删除压缩包

!unzip Stanford3dDataset_v1.2_Aligned_Version.zip  > /dev/null
!rm Stanford3dDataset_v1.2_Aligned_Version.zip

运行点云预处理程序

cd /content/point_edg

!python data_utils/collect_indoor3d_data.py
!rm -r /content/point_edg/data/s3dis/Stanford3dDataset_v1.2_Aligned_Version

cd /content/point_edg

!mv /content/point_edg/data/stanford_indoor3d /content/point_edg/data/s3dis

开始训练

!python train_semseg.py --model point_edg --test_area 5 --log_dir point_edg --batch_size 32 --npoint 4096 --epoch 128

库中自带了PointNet、PointNet++与我用于测试新建的模型PointEdge，若需使用PointNet++就将命令中的point_edg改成pointnet2_sem_seg。
训练log与模型参数会储存在云盘~/PCT_log/sem_seg目录下，以参数--log_dir命名的文件夹中。

测试模型

!python test_semseg.py --log_dir point_edg --test_area 5 --visual

该命令如带了 --visual 参数则会在log文件夹下的visual文件夹里找到该模型分割的点云文件（会将area5全部模型的groundtruth和你模型分割的结果保存，差不多6G大小，小心云盘空间不足）。
将得到的点云文件下载，后缀改成.txt导入CloudCompare软件中查看分割效果。

上图分别为ground_truth与我的PointEdge模型的语义分割效果。