【点云处理之论文狂读经典版8】—— O-CNN: Octree-based Convolutional Neural Networks for 3D Shape Analysis

摘要

• 以八叉树叶子节点中的平均法向量为输入，执行3D CNN操作
• 设计了一种新的八叉树数据结构，并行化程度高，卷积操作更适合GPU进行处理
• 通过限制在八叉树子节点上的计算，O-CNN的内存消耗和计算复杂度随着八叉树深度的增加呈二次增长，使得3D CNN能够处理更高分辨率的3D 模型
• 代码详见：https://github.com/Microsoft/O-CNN

方法

在给定一个带有方向的三维模型后，首先构造八叉树，然会对八叉树中用于CNN操作的信息进行封装。

用于3D CNN的八叉树

八叉树构造

• 将3D模型缩放到一个轴对齐的单位立方体内
• 以宽度优先(breadth-first)的顺序递归地划分3D形状的立方体
  • 在当前深度$l$中遍历所有包含3D形状边界的非空octants
  • 在下一个深度$l+1$ 中将其划分到8个孩子octants中
  • 重复以上步骤直至到达深度$d$， 见图2(a-b)

构造完八叉树后，在每个octant中收集一组用于CNN的属性(properties)，并将这些值存储在八叉树中。具体而言，就是计算八叉树中每个octant的shuffle key和label。与此同时，我们的方法从叶子节点中的3D形状提取CNN的输入信号，并将最终的CNN特征存储在每个octant中。

如图2所示，根据深度顺序整理了存储在八叉树中的数据。在每个深度中，根据shuffle key的升序对octant进行排序，并将它们的属性值分配到一组1D的属性向量中。所有的属性向量都共享相同的索引，向量的长度和当前层的octant数量相同。

shuffle key

对于深度为$l$的一个octant $O$上的shuffle key，使用一个唯一的$3l$比特字符$\mathrm{key}(O):=$ $x_1{y}_1{z}_1{x}_2{y}_2{z}_2\dots x_l{y}_l{z}_l$对其在三维空间中的位置进行编码，其中每一个三比特组 x i , y i , z i ∈ { 0 , 1 } x_{i}, y_{i}, z_{i} \in\{0,1\} xi​,yi​,zi​∈{0,1} 定义了该octant的父亲octant在3D立方体的相对位置。一个octant $O$的整数坐标$(x,y,z)$通过 $x=\left(x_1{x}_2\dots x_l\right),y=$ $\left(y_1{y}_2\dots y_l\right),z=\left(z_1{z}_2\dots z_l\right)$决定。通过对shuffle key采取升序的方式对octant进行排序，并将在第$l$层的所有octant的shuffle key存储在一个shuffle key向量$S_l$中，用于后面构造actant的3D卷积邻域。

label

在实现CNN时，pooling操作被频繁地用于从第$l$层的octant下采样到第$l-1$层的父亲octant中。因此，我们需要在相邻层中快速的找到octants间的父亲-孩子关系。

因此，为第$l$层的非空octant设计了一个label $p$，表明这是在第$l$层排好序的octant列表中第$p$个非空octant。对于空octant，将它们的值设置为0。在第$l$层中所有octant的label都被排序成label向量$L_l$。

给定一个在第$l$层索引为$j$的顶点，通过$k=8\times \left(L_l\left[j\right]-1\right)$计算其在$l+1$层中第一个孩子的索引$k$，这是基于以下两点得到的：

1. 在第$l$层，仅非空的octant才会被划分
2. 由于octants是根据shuffle key的升序进行排序的，因此octant的8个孩子都是按照顺序存储的

此外，孩子octant和它们的非空父亲在属性向量中也遵循相同的顺序。

Input signal

使用在叶子octant中计算得到的平均法向量作为CNN的输入信号。对于空叶子octant，简单地分配零向量作为输入信号；对于非空叶子octant，我们用一组点对嵌入在叶子octant中的3D形状表面进行采样，并将所有采样点的法线取平均值，作为该叶子octant的输入信号。将所有叶子octant的输入信号存储到一个输入信号向量中。该向量的大小是八叉树中叶子octant数量的大小。

与之前许多基于voxel的CNN方法使用过的binary indicator函数不同，法向量信号非常稀疏，只在表面上才会有非零值，在叶子octant上采样的平均法向量能够更好地表示局部3D形状的方向，并且能够为CNN提供更真实的3D形状信息。图3比较了一个被体素化的3D模型和一个用八叉树表示的相同3D模型，其中叶子octant的大小与voxel的大小相同。

CNN features

对于定义在第$l$层上的每个3D卷积核，都在特征图向量$T_l$中记录第$l$层所有octant的卷积结果。

Mini-batch of 3D models

在CNN训练过程中，一个mini-batch中的3D形状对应的八叉树各不相同。为了能够使CNN训练在GPU上更高效，将这些八叉树合并成一个超八叉树(super-octree)。在每个八叉树深度$l$上，拼接所有3D 目标的属性向量($S_l$，$L_l$，$T_l$)，得到超八叉树的属性向量$S_l^{\ast }$，$L_l^{\ast }$，$T_l^{\ast }$。

在此之后，我们通过使用每个shuffle key中最高的8 bit存储对象的索引，从而更新$S_l^{\ast }$中的shuffle key。还更新超八叉树中的label向量$L_l^{\ast }$，用于表示每个非空octant在超八叉树中的索引。

八叉树上的CNN操作

3D Convolution

为了能够在八叉树上应用卷积操作，需要在相同八叉树深度上选择相邻的octants。为了能够更高效地计算卷积，以展开的形式写出卷积操作${\Phi }_{\mathrm{c}}$：
$${\Phi }_{\mathrm{c}}(O)=\sum _n\sum _i\sum _j\sum _k{W}_{ijk}^{(n)}\cdot T^{(n)}\left(O_{ijk}\right).$$
其中$O_{ijk}$表示$O$的相邻octants，$T(\cdot )$ 表示与$O_{ijk}$相关的特征向量。$T^{(n)}(\cdot )$表示特征向量的第$n$维通道，$W_{ijk}^{(n)}$是卷积操作的权重。如果$O_{ijk}$在八叉树中不存在，那就将$T\left(O_{ijk}\right)$设置为零向量。通过这种方式，卷积操作可以转化为矩阵乘积的形式，能够在在GPU上高效地运算。

卷积核大小为$K$的卷积操作需要$K^3-1$个相邻octants，若$K$太大，会导致I/O处理的巨大消耗。

因此构建了哈希表$\mathcal{H}:key(O)\mapsto \mathrm{index}(O)$ 用于更高效地进行搜索，其中 index $(O)$记录着 $O$在 $S$中的位置。由于哈希表的 amortized 时间复杂度是常数，因此这种选择可以被视为CNN计算成本和内存成本的平衡。在给定$S_l$中存储的shuffled key后，octant的 integer coordinates $(x,y,z)$ 可以被存储，接着邻域octant的坐标可以以恒定的时间被计算，得益于其对应的shuffled key。在给定相邻octant的shuffled key后，哈希表能够并行地搜寻它们的索引，根据这些索引收集相邻数据信息，然后应用卷积运算。

如果卷积的stride是1，那么上述操作可以被用于所有八叉树的当前层。正常情况下，对于每个octant，哈希表都会搜寻$K^3-1$次。但是，在同一个父亲节点下8个兄妹octant会有很多的重叠，而且8个octants一共才有$(K+1{)}^3$个单独的octant。可以通过寻找8个兄妹octant所有的$(K+1{)}^3-8$个相邻octants进一步缩短领域搜寻时间。

如果卷积的stride是2，实现起来更简单。对于相同父亲octant的8个兄妹octant，卷积操作可以被用到第一个兄妹octant上，同时忽略掉其他的兄妹octant，等同于对特征图的分辨率执行因子为2的下采样。

如果卷积的stride为$2^r(r>1)$，那么卷积操作就会被用在高度为$r$的每个sub-tree的第一个octant上，那么深度图分辨率将会执行因子为$2^r(r>1)$的下采样。

由于八叉树特殊的分层结构，卷积的stride只能被限制在2的幂的整数。

当执行的stride大于1时，会产生下采样现象，data向量$T$就会被缩短。数据从底层流向上层。那么存储在$L_l$中的信息会被用于获得correspondence，例子见图2 （c）。$T_2$最初的长度为12，下采样后变成了3，但是该层有4个octant，这时就要结合存储在$L_1$中的信息对$T_1$中的信息进行更新。

Pooling

Pooling的主要作用就是不断地压缩表示的空间尺寸。Pooling层独立作用于特征图的每个通道，并改变特征图的大小。最常用的形式便是the max-pooling层，该层的filters大小为2，stride为2。

在八叉树上应用pooling也很方便。由于属于相同父亲octant的8个兄妹octant都是连续存储的，在八叉树上用于pooling操作便是从这8个octant选择最大的元素，能够在GPU上高效地执行。接着特征图的分辨率会执行因子为2 的下采样，父亲octant的信息将会被用于下一步的操作。

实际上，pooling操作可以被看成是一种特殊的卷积形式，对于其他lernel大小或stride大小的一般池操作，我们遵循前面介绍的方法，即找到相应的相邻octant并应用特定的操作，如max-pooling或average pooling。

Unpooling

Unpooling操作是pooling的逆操作，并起到上采样的作用。The max-unpooling操作通常和max-pooling操作一起使用。在应用max-pooling操作后，每个pooling区域内的最大值位置会被记录成一组switch 变量，并将这些变量存储到一个连续的阵列。The max-unpooling操作会使用这些swithes，将当前特征图的信号放到上采样特征图中最合适的位置。

Deconvolution

反卷积操作被用于增加特征图的密度，可以通过逆向进行卷积的前馈和反馈传递实现该操作。

Remark

本文的CNN操作仅作用在octants上，所以就没必要在空区域中传播信息，也不会在空区域中交换信息。通过限制信息在八叉树中的传播，形状信息就会沿着形状更高效地进行交换。

通过限制octant中的数据存储和CNN计算，基于CNN的八叉树的计算复杂度为$\mathcal{O}(n^2)$，其中$n$是在叶子节点层上每个维度的voxel分辨率。全voxel方法的计算复杂度为$\mathcal{O}(n^3)$。此外，由于内存中存储的数据是连续的，O-CNN与定义在规则网格上2D CNN和3D CNN 的GPU计算性能一样高。在训练时，第$k$次迭代所需要的哈希表和相邻信息可以在第$(k-1)$次迭代中提前计算好。

网络结构

为了更清晰的证明本文提出的基于八叉树表示方法具有优势，遵循LeNet的概念设计了一个简单的网络。

记${\mathrm{U}}_l$ = convolution + BN + ReLU + pooling为一个基本单元，卷积作用在第$l$层的octants上。对于基本单元${\mathrm{U}}_l$，特征图的通道数量设置为$2^{\mathrm{m}\mathrm{a}\mathrm{x}(l,9-l)}$，卷积kernel大小为3，O-CNN被定义为：
$$\mathrm{i}\mathrm{n}\mathrm{p}\mathrm{u}\mathrm{t}\to {\mathrm{U}}_d\to {\mathrm{U}}_{d-1}\to \cdots \to {\mathrm{U}}_2$$
将上述网络记为O-CNN(d)。为了从不同地八叉树结构中对齐所有特征，强制将所有的第二层octants存在，并使用零向量填充第二层所有地空octant。

O-CNN for shape analysis

对于object classification，采用下面的结构：
$$\mathrm{O}-\mathrm{C}\mathrm{N}\mathrm{N}(\mathrm{d})\to \mathrm{D}\mathrm{r}\mathrm{o}\mathrm{p}\mathrm{o}\mathrm{u}\mathrm{t}\to \mathrm{F}\mathrm{C}(128)\to \mathrm{D}\mathrm{r}\mathrm{o}\mathrm{p}\mathrm{o}\mathrm{u}\mathrm{t}\to \mathrm{F}\mathrm{C}({\mathrm{N}}_{\mathrm{c}})\to \mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\to \mathrm{o}\mathrm{u}\mathrm{t}\mathrm{p}\mathrm{u}\mathrm{t}$$

对于shape retrieval，object classification的输出作为key来搜寻与要寻找目标最相似的形状。

对于shape part segmentation，利用最新的网络DeconvNet，为了更好的预测，在卷积网之后还级联一个反卷积网络。卷积网络可以设置为O-CNN(d)，反卷积网络是O-CNN(d)的镜像，其中卷积网络和pooling操作被替换成了反卷积和unpooling操作。
记${\mathrm{D}\mathrm{U}}_l$=unpooling + deconvolution +BN + ReLU为基本单元，shape segmentation 的网络结构可以表示为：

$$\mathrm{O}-\mathrm{C}\mathrm{N}\mathrm{N}(\mathrm{d})\to {\mathrm{D}\mathrm{U}}_2\to {\mathrm{D}\mathrm{U}}_3\to \cdots \to {\mathrm{D}\mathrm{U}}_d$$

实验

Object classification

Shape retrieval

展望

• 适应性八叉树. 对于一些表面平整的区域，没必要再对八叉树进行划分，而使用一个大的octant就足够了，这样可以节省更多的内存
• 通用的lattices.
• 网络结构. 残差网络和循环神经网络前景很好