当前位置：首页 > news >正文

LitePT：轻量级点云Transformer架构设计与优化

news 2026/5/2 3:36:43

1. 项目概述

LitePT是一个专为点云数据处理设计的轻量级Transformer架构。在计算机视觉领域，点云作为三维空间数据的重要表示形式，其处理一直面临着独特的挑战。传统点云处理方法往往需要复杂的特征工程，而基于Transformer的架构虽然表现出色，但计算开销大、内存占用高的问题限制了其在移动端和边缘设备上的应用。

这个架构最吸引我的地方在于它巧妙地在模型性能和计算效率之间找到了平衡点。通过一系列创新设计，LitePT在保持Transformer强大表征能力的同时，将参数量和计算复杂度降低到了传统方法的1/3左右。在实际测试中，它在ModelNet40数据集上达到了92.3%的分类准确率，推理速度比标准Point Transformer快了近5倍。

2. 核心设计思路拆解

2.1 点云数据的特殊性处理

点云数据与常规图像数据有本质区别：无序性、稀疏性和非均匀性。LitePT采用了几项关键设计来应对这些特性：

局部-全局注意力机制：在3D空间中将点云划分为局部区域，先在局部范围内计算注意力，再通过层次化方式聚合全局信息。这种设计避免了传统Transformer在点云上直接计算全局注意力的高计算成本。
位置编码优化：采用可学习的相对位置编码，而非固定的正弦编码。具体实现是通过MLP学习每个点与其k近邻点的相对位置关系，公式表示为：
```
PE_ij = MLP(p_i - p_j)
```
其中p_i和p_j是点的坐标。

2.2 轻量化技术实现

2.2.1 注意力矩阵压缩

标准Transformer的注意力矩阵大小为N×N（N为点数），这在点云场景下（通常N>1000）会带来巨大计算负担。LitePT采用两种压缩策略：

关键点采样：使用最远点采样(FPS)选取M个关键点(M<<N)，只在关键点间计算注意力
低秩近似：将注意力矩阵分解为两个低秩矩阵的乘积，维度从N×N降为N×k和k×N(k≈32)

2.2.2 混合精度计算

在模型实现中，对不同的计算模块采用不同的精度：

注意力计算：FP16
特征变换：FP32
残差连接：FP16

这种混合精度策略在保持数值稳定性的同时，减少了约40%的内存占用。

3. 架构实现细节

3.1 网络整体结构

LitePT采用U-Net式的编码器-解码器结构：

输入点云 → 下采样模块(4×) → LitePT模块(3层) → 上采样模块 → 分类/分割头

每个LitePT模块包含：

局部注意力层（4头）
轻量MLP（隐藏层维度=64）
层归一化
残差连接

3.2 关键代码实现

以局部注意力计算为例：

class LocalAttention(nn.Module): def __init__(self, dim, num_heads=4, k=16): super().__init__() self.num_heads = num_heads self.k = k self.qkv = nn.Linear(dim, dim*3) self.proj = nn.Linear(dim, dim) def forward(self, x, pos): B, N, C = x.shape # 获取k近邻 idx = knn(pos, k=self.k) # (B, N, k) # 计算query/key/value qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C//self.num_heads) q, k, v = qkv.unbind(2) # (B, N, h, C/h) # 局部注意力计算 k = index_points(k, idx) # (B, N, k, h, C/h) v = index_points(v, idx) attn = (q.unsqueeze(2) @ k.transpose(-2,-1)) # (B, N, k, h) attn = attn / (C//self.num_heads)**0.5 attn = attn.softmax(dim=2) x = (attn.unsqueeze(-1) * v).sum(dim=2) # (B, N, h, C/h) x = x.transpose(1,2).reshape(B, N, C) return self.proj(x)