- 论文题目:Point Transformer
- 发布期刊:ICCV
- 通讯地址:牛津大学 & 香港大学 & 香港中文大学
- 代码地址:https://github.com/POSTECH-CVLab/point-transformer
介绍
- 提出了Point Transformer层:基于自注意力机制,设计了适用于3D点云处理的自注意力层。由于点云本质上是嵌入到3D空间中的点集,自注意力机制在这种情况下很自然地适用。该层对点的排列顺序不敏感,适用于3D点云数据。
- 用于多种3D理解任务:作者通过Point Transformer网络,处理语义场景分割、物体部分分割和物体分类等任务,展示了模型在不同领域的广泛适用性。实验表明,Point Transformer在多个基准数据集上取得了最新的最佳性能。
- 实验结果:在S3DIS大规模语义分割数据集上,Point Transformer的mIoU达到70.4%,首次突破了70%的门槛。此外,模型在3D形状分类和物体部分分割等任务上也表现优异,分别在ModelNet40和ShapeNetPart数据集上取得了93.7%和86.6%的准确率。
- 创新点:
- 提出了高效的局部自注意力机制,使模型可以扩展到处理包含数百万点的大规模3D场景。
- 引入了位置编码,在自注意力生成和特征转换分支中使用相对位置编码,显著提升了模型的性能。
核心思想
Point Transformer的核心思想和原理是将Transformer架构引入到3D点云处理任务中。由于3D点云数据的无序性和稀疏性,传统的卷积神经网络(CNN)无法直接应用于点云数据。Point Transformer通过自注意力机制,克服了这一挑战,能够灵活处理不规则的点集。
-
自注意力机制(Self-Attention Mechanism)适合点云数据:
- 点云数据本质上是无序且嵌入到三维空间中的点集,传统卷积操作依赖于规则的像素网格,因此难以直接应用于点云数据。Transformer中的自注意力机制是集合运算,不依赖于输入数据的排列顺序和大小,能够自然地处理无序点云。
- 自注意力可以捕捉每个点与其邻域内其他点之间的关系,使模型可以灵活地聚合点云数据中的局部信息。
-
局部自注意力(Local Self-Attention):
- 点云数据的规模通常很大,直接在全局范围内应用自注意力会产生巨大的计算开销。Point Transformer通过局部自注意力机制,仅在每个点的局部邻域内应用自注意力,从而降低计算复杂度。邻域通常通过K近邻算法(kNN) 来选择。
- 这种局部自注意力的方式确保了模型能够扩展到处理大规模点云场景,而不会因为数据规模的增长而计算量爆炸。
-
位置编码(Positional Encoding)增强几何结构理解:
- 位置编码在自注意力模型中扮演重要角色,帮助模型捕捉点云中各点之间的空间关系。Point Transformer引入了可训练的相对位置编码,将点的三维坐标(x, y, z)直接作为输入,增加了自注意力机制对点之间几何结构的敏感性。
- 位置编码不仅影响注意力权重的计算,还作用于特征变换过程中,使得模型能够根据点与点之间的相对位置,生成更加准确的几何特征表示。
-
向量自注意力(Vector Attention)更灵活的特征聚合:
- 与传统的标量自注意力不同,Point Transformer引入了向量自注意力机制。标量自注意力将所有特征通道使用相同的注意力权重进行聚合,而向量自注意力则允许不同特征通道使用独立的注意力权重,从而提升模型的表达能力。
- 向量自注意力能够对每个特征通道进行更细粒度的调制,从而更好地聚合点云的局部信息。
Point Transformer的原理基于将自注意力运算应用于3D点云数据中的每个点。主要步骤包括:
-
输入特征处理:
- 对于每个输入点,首先通过一个多层感知器(MLP)将其初始特征映射到高维空间。这一步通过点的三维坐标或其他附加信息(如颜色、法线等)生成初始特征向量。
-
局部自注意力操作:
- 对于每个点,使用K近邻(kNN)选择该点的邻域(通常为16个邻居点),在该局部邻域内计算自注意力。自注意力的核心是根据相邻点的特征和位置生成注意力权重,并使用这些权重对邻域内的点特征进行加权求和。
-
位置编码的作用:
- 位置编码函数 (\delta = \theta(p_i - p_j)),其中 (p_i) 和 (p_j) 是点的三维坐标,(\theta) 是一个多层感知器(MLP)。位置编码用于在特征变换和自注意力权重计算中反映点与点之间的相对位置,确保模型能够捕捉到点云中的几何信息。
-
残差连接:
- 每个Point Transformer层都包括残差连接,使得模型能够有效地训练更深的网络,缓解梯度消失问题。
-
降采样与上采样:
- 在网络的编码器部分,使用降采样模块(如最远点采样)逐步减少点的数量,同时使用kNN聚合局部特征,提取更高层次的几何特征。
- 在解码器部分,使用上采样模块(如三线性插值)恢复点的空间分辨率,同时结合多层特征信息,生成最终的输出(如分割结果或分类结果)。
Point Transformer的核心思想是通过局部自注意力机制处理3D点云数据,同时引入位置编码和向量自注意力机制来增强对点云几何结构的理解。通过这一方法,Point Transformer能够在3D点云分类、语义分割、物体部分分割等任务中表现出色,超越了传统的卷积神经网络和其他基于图或卷积的方法。
实现原理和代码实现
Point Transformer的整体网络结构基于编码器-解码器(Encoder-Decoder)架构,主要包括以下几个关键模块:输入预处理、Point Transformer块(核心模块)、Transition Down(降采样)、Transition Up(上采样)以及输出头。它主要通过局部自注意力机制来处理点云数据,逐步提取高层次的特征,最后用于分类、分割等3D点云任务。
Point Transformer Block(核心模块)
- 局部自注意力机制:每个Point Transformer Block都应用局部自注意力机制。基于K近邻选择每个点的局部邻域点,利用自注意力机制来计算点与其邻域点之间的关系。自注意力是整个网络中的主要特征聚合方式,能够捕捉点云数据的局部几何结构。
- 残差连接:类似于ResNet,每个Point Transformer Block包含残差连接,确保深层网络的信息传递和梯度稳定。
- 位置编码:位置编码将每个点的三维坐标信息(位置)引入到自注意力计算中。它不仅用于生成注意力权重,还在特征变换中使用,帮助模型理解点之间的空间关系。
class PointTransformerBlock(nn.Module):
def forward(self, pxo):
identity = x # 保留输入用于残差连接
x = self.relu(self.bn1(self.linear1(x))) # 线性变换+归一化+激活
x = self.relu(self.bn2(self.transformer2([p, x, o]))) # 应用PointTransformerLayer进行自注意力操作
x = self.bn3(self.linear3(x)) # 最后的线性变换和归一化
x += identity # 残差连接
x = self.relu(x)
return [p, x, o]
- 功能:在Point Transformer中引入残差连接,类似于ResNet的残差块结构。
- 实现:
- 输入特征先经过线性变换、批归一化和ReLU激活,再通过
PointTransformerLayer
进行自注意力运算。 - 最后通过残差连接,将原始输入特征与处理后的特征相加,再进行激活。
- 输入特征先经过线性变换、批归一化和ReLU激活,再通过
- 输出:残差连接后的点特征。
class PointTransformerLayer(nn.Module):
...
def forward(self, pxo) -> torch.Tensor:
p, x, o = pxo # p: (n, 3) 点坐标, x: (n, c) 特征, o: (b) 批次大小
x_q, x_k, x_v = self.linear_q(x), self.linear_k(x), self.linear_v(x) # 计算查询(query)、键(key)、值(value)
x_k = pointops.queryandgroup(self.nsample, p, p, x_k, None, o, o, use_xyz=True) # 提取局部邻域点的特征
x_v = pointops.queryandgroup(self.nsample, p, p, x_v, None, o, o, use_xyz=False)
...
w = self.softmax(w) # 计算注意力权重
x = ((x_v + p_r) * w.unsqueeze(2)).sum(1).view(n, c) # 聚合局部邻域特征
return x
- 功能:这是局部自注意力机制的实现。主要作用是通过自注意力机制聚合点云的局部特征。
- 输入:
p
:点的3D坐标。x
:点的特征。o
:批次大小。
- 实现:
linear_q
,linear_k
,linear_v
计算点的查询(query)、键(key) 和 值(value)。- 使用
pointops.queryandgroup
函数提取每个点的局部邻域点(K近邻点)的特征。 - 通过点与其邻域点之间的差异计算注意力权重,再对邻域点的特征进行加权求和,生成更新的特征。
- 输出:返回聚合后的特征。
Transition Down(降采样模块)
- 降采样操作:降采样模块通过最远点采样(Farthest Point Sampling, FPS)从当前点集中选取子集,这样可以减少点的数量,同时保留全局空间中的代表性点。
- 局部聚合:降采样之后,使用K近邻聚合每个点的局部特征,并将特征映射到低维空间。这有助于减少计算开销,同时确保模型的上下文信息完好无损。
class TransitionDown(nn.Module):
def forward(self, pxo):
p, x, o = pxo # 输入点的坐标、特征和批次信息
if self.stride != 1:
idx = pointops.furthestsampling(p, o, n_o) # 进行最远点采样,降采样点云
x = pointops.queryandgroup(self.nsample, p, n_p, x, None, o, n_o, use_xyz=True) # 提取邻域特征
x = self.relu(self.bn(self.linear(x).transpose(1, 2).contiguous())) # 线性变换+批量归一化
x = self.pool(x).squeeze(-1) # 池化
else:
x = self.relu(self.bn(self.linear(x)))
return [p, x, o]
- 功能:降采样点云,类似于卷积神经网络中的池化操作。
- 实现:
- 如果
stride != 1
,则进行最远点采样(Furthest Point Sampling, FPS),减少点的数量。 - 通过
queryandgroup
提取降采样后每个点的局部邻域特征,之后对特征进行线性变换、批归一化和池化。
- 如果
- 输出:降采样后的点和特征。
Transition Up(上采样模块)
- 上采样操作:上采样模块恢复降采样过程中丢失的空间分辨率。通过三线性插值(Trilinear Interpolation)操作,将低分辨率的点集特征映射回高分辨率点集。
- 特征融合:在上采样过程中,使用跳跃连接(Skip Connection)将编码阶段每一层的特征与解码阶段的特征结合。这种方式类似于U-Net架构,确保模型在恢复分辨率的同时能够保留更多的细粒度信息。
class TransitionUp(nn.Module):
def forward(self, pxo1, pxo2=None):
if pxo2 is None:
# 如果没有额外输入,则直接对输入进行特征聚合和线性变换
x_tmp = []
for i in range(o.shape[0]):
x_b = torch.cat((x_b, self.linear2(x_b.sum(0, True) / cnt).repeat(cnt, 1)), 1) # 聚合特征
x = self.linear1(x)
else:
# 如果有来自上一级的输入,进行插值上采样
x = self.linear1(x1) + pointops.interpolation(p2, p1, self.linear2(x2), o2, o1)
return x
- 功能:上采样点云,恢复空间分辨率,类似于卷积神经网络中的上采样操作。
- 实现:
pxo2
为上一级点云。如果没有pxo2
,对当前特征直接进行线性变换和特征聚合。- 如果有
pxo2
,通过插值操作将较低分辨率的特征插值到高分辨率的点云上。
- 输出:上采样后的点特征。
如何改进PointNet++
PointNet++ 是一种用于3D点云处理的深度学习模型,通过递归地应用 PointNet 在局部邻域内提取特征,但它依赖于固定的几何关系和简单的特征聚合方法。相比之下,Point Transformer 引入了更强大的 自注意力机制,允许模型更灵活地捕捉局部和全局信息。因此,将 Point Transformer 的思想应用于 PointNet++ 可以进一步提升点云处理的能力。
以下是三种将 Point Transformer 的思想应用到 PointNet++ 的改进方案:
1. 引入局部自注意力机制替代固定邻域特征聚合
- 改进方式:PointNet++ 中的局部特征聚合是通过基于几何距离的简单 max pooling 或 average pooling 操作完成的。可以通过 Point Transformer 的局部自注意力机制 替代这种固定的特征聚合方式。
- 原理:自注意力机制能够在邻域点之间灵活地计算注意力权重,而不仅仅是依赖固定的几何距离。这意味着特征可以根据每个点之间的相似性自适应地聚合,而不是仅仅考虑它们之间的物理距离。
- 好处:
- 可以更智能地捕捉点与点之间的关系,特别是在复杂几何形状中。
- 增强模型对局部结构的感知,适应性更强。
- 实现思路:
- 在 PointNet++ 的每个局部邻域特征提取阶段(即
set abstraction layer
),引入 Point Transformer 中的自注意力机制,替代原有的 max/average pooling 操作,使得特征根据注意力机制动态加权。
- 在 PointNet++ 的每个局部邻域特征提取阶段(即
2. 使用位置编码增强几何信息的表达
- 改进方式:PointNet++ 仅使用几何坐标作为输入特征,并在局部邻域中通过直接计算点与点之间的距离来聚合信息。可以通过 Point Transformer 中的可训练位置编码 来增强对点与点之间相对几何位置的捕捉。
- 原理:Point Transformer 中使用了一个多层感知器(MLP)来为点的几何坐标生成位置编码,进一步提升了模型对几何结构的理解能力。这可以弥补 PointNet++ 中几何信息表达的不足。
- 好处:
- 增加对点之间相对位置的敏感度,使得模型能够更准确地捕捉复杂的几何关系。
- 提升模型在具有细粒度几何结构的场景中的表现。
- 实现思路:
- 在 PointNet++ 的
set abstraction layer
中引入一个 MLP 来生成每个点的相对位置编码,并将其与现有特征结合,通过 Point Transformer 的方式进行特征增强。
- 在 PointNet++ 的
3. 应用全局自注意力机制改进全局特征学习
- 改进方式:PointNet++ 中对全局特征的学习是通过多次局部特征聚合逐层递进的,可以通过 Point Transformer 的全局自注意力机制 来提升全局特征学习的能力。
- 原理:PointNet++ 在提取全局特征时依赖逐层的特征汇聚(例如使用全局 max pooling),而 Point Transformer 的自注意力机制可以在不受层级限制的情况下直接捕捉全局特征。通过在全局层面引入自注意力机制,模型可以更好地捕捉点之间的长距离依赖关系。
- 好处:
- 提升模型在需要理解全局上下文信息(例如大型场景、复杂物体)的表现。
- 增强模型对全局特征的表达能力,特别是在复杂场景中。
- 实现思路:
- 在 PointNet++ 的
global feature aggregation
阶段,引入 Point Transformer 的全局自注意力机制,替代原有的全局 max pooling,使得模型能够通过自注意力来捕捉全局上下文信息。
- 在 PointNet++ 的