点云中的Minkowski卷积

一、MinkowskiEngine简介

第一次见到MinkowskiEngine，应该是在两个月之前了，当时也没有去留意这个库。最近读了一些点云的论文，发现还是有不少论文的源码是基于MinkowskiEngine的，包括PointContrast(ECCV 2020)，DGR(Deep Global Registration, CVPR 2020)，Learning Multiview 3D Point Cloud Registration(CVPR 2020)和FCGF(Fully Convolutional Geometric Features, ICCV 2019)等，所以决定去了解一下这个库。

MinkowskiEngine是稀疏张量的自动微分库，致力于高维空间稀疏数据的操作。它支持所有的神经网络层，如Conv, Pooling, Unpooling和广播操作。基于MinkowskiEngine，可以实现点云的分割、分类、重建、补全、检测等任务。

MinkowskiEngine是在4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks[CVPR 2019]被提出的，主页，源码。目前已更新到v0.5版本，且正处于更新中。

二、Minkowski Convolution的定义

点云数据的表示

MinkowskiEngine把点云表示成两部分: 坐标矩阵 C C C和特征矩阵 F F F。

C = [ x 1 y 1 z 1 b 1 : : x N y N z N b N ] C = \left[ \begin{matrix} x_1 & y_1 & z_1 & b_1\\ & : & : & \\ x_N & y_N & z_N & b_N \end{matrix} \right] C=⎣⎡x1xNy1:yNz1:zNb1bN⎦⎤ 和 F = [ f 1 T : f N T ] F = \left[ \begin{matrix} f_1^T \\ : \\ f_N^T \end{matrix} \right] F=⎣⎡f1T:fNT⎦⎤

其中 ( x i , y i , z i ) (x_i, y_i, z_i) (xi,yi,zi)表示点云的坐标， b i b_i bi表示 ( x i , y i , z i ) (x_i, y_i, z_i) (xi,yi,zi)属于batch中的哪个点云(MinkowskiEngine也是把点云组织成batch进行训练)， N N N表示1个batch中所有点的数量， f i T f_i^T fiT表示第 i i i个点的特征，可以是1维或者3维或者其它维度的。

这样的表示，相比于3D卷积(X, Y, Z, D)的表示，可以节省空间。[N << XYZ -> N * 4 + N*D << XYZD, << 表示远小于]
常规3D卷积:

x u out = Σ i ∈ V ( K ) W i x u + i in , for u ∈ Z 3 \text{x}_u^{\text{out}} = \Sigma_{\text{i} \in V(K)}W_i\text{x}_{u + \text{i}}^{\text{in}}, \quad \text{for} \quad u \in \mathbb Z^3 xuout=Σi∈V(K)Wixu+iin,foru∈Z3

u ∈ Z 3 u \in \mathbb Z^3 u∈Z3表示3D坐标， K K K表示卷积中的kernel_size，V(K)是3维空间中的offsets集合， W i ∈ N out × N in W_i \in \mathbb N^{\text {out}} \times \mathbb N^ \text{in} Wi∈Nout×Nin
Minkowski 卷积

x u out = Σ i ∈ N ( u , C in ) W i x u + i in , for u ∈ C out \text{x}_u^{\text{out}} = \Sigma_{\text{i} \in N(u, \mathbb C^{\text{in}})}W_i\text{x}_{u + \text{i}}^{\text{in}}, \quad \text{for} \quad u \in \mathbb C^{\text{out}} xuout=Σi∈N(u,Cin)Wixu+iin,foru∈Cout

对于常规3D卷积，可以看到变化的是 u ∈ C out u \in \mathbb C^{\text{out}} u∈Cout和 i ∈ N ( u , C in ) \text{i} \in N(u, \mathbb C^{\text{in}}) i∈N(u,Cin)。 C in \mathbb C^{\text{in}} Cin和 C out \mathbb C^{\text{out}} Cout是预定义的稀疏张量的输入坐标和输出坐标， N ( u , C in ) = { i ∣ u + i ∈ C in , i ∈ V ( K ) } N(u, \mathbb C^{\text{in}}) = \lbrace \text{i} | u + \text{i} \in \mathbb C^{\text{in}}, i \in V(K) \rbrace N(u,Cin)={i∣u+i∈Cin,i∈V(K)}。因此，相比于常规卷积，不是每一个(x, y, z)位置都会有一个卷积的输出，同时并不是每一个offset位置都会参与计算卷积。

三、从使用的角度看MinkowskiEngine

卷积中常见的操作包括Conv, BN, Pooling, FC, transposed Conv等，本章节基于两个很简单的分类网络和分割网络在MinkowskiEngine环境实验上述操作。实验的环境是Ubuntu14.4, Cuda10.2, PyTorch 1.5, Python 3.7, MinkowskiEngine v0.5(下面的实验代码对于其它环境或许也同样支持)。

3.0 网络的输入数据

为了方便观察数据，设batch size=2，第一个点云P1中具有10个点，第二个点云P2中具有6个点，每个点的特征是3维的。下面代码生成P1和P2点云，并转换成MinkowskiEngine的输入数据格式。

import numpy as np
import torch.nn as nn
import MinkowskiEngine as MEdef print_title(s, data):print('='*20, s, '='*20)print(data)if __name__ == '__main__':origin_pc1 = 100 * np.random.uniform(0, 1, (10, 3))feat1 = np.ones((10, 3), dtype=np.float32)origin_pc2 = 100 * np.random.uniform(0, 1, (6, 3))feat2 = np.ones((6, 3), dtype=np.float32)print_title('origin_pc1', origin_pc1)print_title('origin_pc2', origin_pc2)coords, feats = ME.utils.sparse_collate([origin_pc1, origin_pc2], [feat1, feat2])print_title('coords', coords)print_title('feats', feats)input = ME.SparseTensor(feats, coordinates=coords)print_title('input', input)

上述程序的输出:

==================== origin_pc1 ====================
[[83.28334147 28.87414665 44.48401738][43.04924052 34.66068275 28.201644  ][22.51394645 53.53203799 25.68239097][11.39696393 27.68488056 18.02263419][68.04944494 78.4874799  33.54077384][83.11021987 95.29080943 72.42599245][68.96104764 64.38640545 56.64488121][61.26343854 35.13968286 10.67545387][95.5847873  56.20865881  5.97082126][63.43547357 75.31685552 67.71327187]]
==================== origin_pc2 ====================
[[21.01681082 32.60864402 14.68910937][76.90920828 40.72511594 17.21551445][67.84378491 80.58219012 43.75387818][45.97922404 77.97593435  3.17289328][39.91144138 80.02990713 44.97847053][ 1.55805162 57.33833007 92.04541106]]
==================== coords ====================
tensor([[ 0, 83, 28, 44],[ 0, 43, 34, 28],[ 0, 22, 53, 25],[ 0, 11, 27, 18],[ 0, 68, 78, 33],[ 0, 83, 95, 72],[ 0, 68, 64, 56],[ 0, 61, 35, 10],[ 0, 95, 56,  5],[ 0, 63, 75, 67],[ 1, 21, 32, 14],[ 1, 76, 40, 17],[ 1, 67, 80, 43],[ 1, 45, 77,  3],[ 1, 39, 80, 44],[ 1,  1, 57, 92]], dtype=torch.int32)
==================== feats ====================
tensor([[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.]])
==================== input ====================
SparseTensor(coordinates=tensor([[ 0, 83, 28, 44],[ 0, 43, 34, 28],[ 0, 22, 53, 25],[ 0, 11, 27, 18],[ 0, 68, 78, 33],[ 0, 83, 95, 72],[ 0, 68, 64, 56],[ 0, 61, 35, 10],[ 0, 95, 56,  5],[ 0, 63, 75, 67],[ 1, 21, 32, 14],[ 1, 76, 40, 17],[ 1, 67, 80, 43],[ 1, 45, 77,  3],[ 1, 39, 80, 44],[ 1,  1, 57, 92]], dtype=torch.int32)features=tensor([[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.],[1., 1., 1.]])coordinate_map_key=coordinate map key:[1, 1, 1]coordinate_manager=CoordinateMapManagerCPU([1, 1, 1]:    CoordinateMapCPU:16x4algorithm=MinkowskiAlgorithm.DEFAULT)spatial dimension=3)

可以看到:

ME.utils.sparse_collate把P1和P2点云数据的坐标进行了量化，并组成了batch的格式, 0表示属于P1点云数据，1表示属于P2点云数据
ME.SparseTensor把数据转换成了SparseTensor，MinkowskiEngine需要的数据格式。SparseTensor包括coordinates和features的信息。

3.1 分类网络

这里只实现Conv(3, 64) + BN + ReLU + GlobalPooling + FC(64, 32)的简单分类网络。

import numpy as np
import torch.nn as nn
import MinkowskiEngine as MEclass ExampleNetwork(ME.MinkowskiNetwork):def __init__(self, in_feat, out_feat, D=3):super(ExampleNetwork, self).__init__(D)self.conv = nn.Sequential(ME.MinkowskiConvolution(in_channels=in_feat,out_channels=64,kernel_size=3,stride=1,dilation=1,bias=False,dimension=D),ME.MinkowskiBatchNorm(64),ME.MinkowskiReLU())self.pooling = ME.MinkowskiGlobalPooling(ME.PoolingMode.GLOBAL_AVG_POOLING_KERNEL)self.linear = ME.MinkowskiLinear(64, out_feat)def forward(self, x):out = self.conv(x)print('conv: ', out.coordinates.size(), out.features.size())out = self.pooling(out)print('pooling: ', out.coordinates.size(), out.features.size())out = self.linear(out)print('linear: ', out.coordinates.size(), out.features.size())return outif __name__ == '__main__':origin_pc1 = 100 * np.random.uniform(0, 1, (10, 3))feat1 = np.ones((10, 3), dtype=np.float32)origin_pc2 = 100 * np.random.uniform(0, 1, (6, 3))feat2 = np.ones((6, 3), dtype=np.float32)coords, feats = ME.utils.sparse_collate([origin_pc1, origin_pc2], [feat1, feat2])input = ME.SparseTensor(feats, coordinates=coords)net = ExampleNetwork(in_feat=3, out_feat=32)output = net(input)for k, v in net.named_parameters():print(k, v.size())

程序运行结果如下:

conv:  torch.Size([16, 4]) torch.Size([16, 64])
pooling:  torch.Size([2, 4]) torch.Size([2, 64])
linear:  torch.Size([2, 4]) torch.Size([2, 32])
conv.0.kernel torch.Size([27, 3, 64])
conv.1.bn.weight torch.Size([64])
conv.1.bn.bias torch.Size([64])
linear.linear.weight torch.Size([32, 64])
linear.linear.bias torch.Size([32])

从上面可以看到:

网络中间层的输出和预期是一样的(要注意的是，pooling 之后的size变成(2, 4)和(2, 64)，是因为输入的数据的batchsize=2)
查看网络的参数也可以通过PyTorch中named_parameters()进行遍历，卷积层的参数的size略有不同，PyTorch中应该是(3, 3, 3, 3, 64)的格式，MinkowskiEngine直接变成了(27, 3, 64)。

3.2 分割网络

输入是两个点云P1, P2，分别具有100个点和6个点，网络的经过Conv(3, 64, stride=2) + BN + ReLU + transposed Conv(64, 4)，结构代码如下所示:

import numpy as np
import torch
import torch.nn as nn
import MinkowskiEngine as ME
import MinkowskiEngine.MinkowskiFunctional as MEFclass ExampleNetwork(ME.MinkowskiNetwork):def __init__(self, in_feat, out_feat, D=3):super(ExampleNetwork, self).__init__(D)self.conv =  ME.MinkowskiConvolution(in_channels=in_feat,out_channels=64,kernel_size=3,stride=2,dilation=1,bias=False,dimension=D)self.bn = ME.MinkowskiBatchNorm(64)self.conv_tr = ME.MinkowskiConvolutionTranspose(in_channels=64,out_channels=4,kernel_size=3,stride=2,dilation=1,bias=False,dimension=D)def forward(self, x):print('input: ', x.coordinates.size(), x.features.size())out = self.conv(x)print('conv: ', out.coordinates.size(), out.features.size())out = self.bn(out)print('bn: ', out.coordinates.size(), out.features.size())out = MEF.relu(out)print('relu: ', out.coordinates.size(), out.features.size())out = self.conv_tr(out)print('conv_tr', out.coordinates.size(), out.features.size())return outif __name__ == '__main__':origin_pc1 = 5 * np.random.uniform(0, 1, (100, 3))feat1 = np.ones((100, 3), dtype=np.float32)origin_pc2 = 100 * np.random.uniform(0, 1, (6, 3))feat2 = np.ones((6, 3), dtype=np.float32)coords, feats = ME.utils.sparse_collate([origin_pc1, origin_pc2], [feat1, feat2])input = ME.SparseTensor(feats, coordinates=coords)net = ExampleNetwork(in_feat=3, out_feat=32)output = net(input)print(torch.equal(input.coordinates, output.coordinates))print(torch.equal(input.features, output.features))

输出结果为:

input:  torch.Size([74, 4]) torch.Size([74, 3])
conv:  torch.Size([31, 4]) torch.Size([31, 64])
bn:  torch.Size([31, 4]) torch.Size([31, 64])
relu:  torch.Size([31, 4]) torch.Size([31, 64])
conv_tr torch.Size([74, 4]) torch.Size([74, 4])
True
False

通过实验可以观察得到：

原始的点云数量应该是100 + 6个，经过量化后输出网络的点云数量是74个。
stride=2的卷积使得点云的数量从74->31个，可以理解为降低了点云的分辨率(对比图像)
transposed Conv把31个点的点云又重新恢复到74个点，倒数第二个True，表示input和output的坐标是一致的，也就是说通过conv + tr_conv，点云的数量和在tensor中的顺序并不会改变。

四、总结

MinkowskiEngine可以以较低显存实现3D卷积的操作
MinkowskiEngine使用方便，Conv, BN, ReLU, transposed Conv等的调用方式同PyTorch，个人感觉可以轻松实现分类、分割等网络架构。
MinkowskiEngine在把数据输入网络之前，会对数据进行量化，会导致部分点云数据的信息丢失。

初次接触MinkowskiEngine，理解不对的地方欢迎大家指正。