【三维点云】Point Transformer核心代码详细注释（自学记录）

        Point Transformer是一种基于注意力机制的点云处理方法。给定一个输入点云，Point Transformer将点云中每个点的位置信息和特征信息作为输入，然后通过一个带有多头注意力机制的Transformer模型来学习一个新的点云表示。

具体实现：

class TransformerBlock(nn.Module):
    def __init__(self, d_points, d_model, k) -> None:
        super().__init__()
        self.fc1 = nn.Linear(d_points, d_model)  # 线性变换，将每个点特征从d_points维转换为d_model维
        self.fc2 = nn.Linear(d_model, d_points)  # 线性变换，将加权求和后的结果从d_model维转换为d_points维
        self.fc_delta = nn.Sequential(  # 学习每个邻居点与中心点之间的位置关系（相对距离）
            nn.Linear(3, d_model),
            nn.ReLU(),
            nn.Linear(d_model, d_model)
        )
        self.fc_gamma = nn.Sequential(  # 神经网络结构，用于学习每个点的权重
            nn.Linear(d_model, d_model),
            nn.ReLU(),
            nn.Linear(d_model, d_model)
        )
        self.w_qs = nn.Linear(d_model, d_model, bias=False)  # 线性变换，生成查询向量
        self.w_ks = nn.Linear(d_model, d_model, bias=False)  # 线性变换，生成邻居点向量
        self.w_vs = nn.Linear(d_model, d_model, bias=False)  # 线性变换，生成邻居点特征向量
        self.k = k

    def forward(self, xyz, features):
        dists = square_distance(xyz, xyz)  # 计算每两个点之间的距离，用于选取每个点最近的K个邻居点
        knn_idx = dists.argsort()[:, :, :self.k]  # 将每个点的邻居点按照距离排序，选取最近的K个邻居点
        knn_xyz = index_points(xyz, knn_idx)  # 获取每个点的K个邻居点的坐标

        pre = features  # 记录一下之前的特征向量，用于后续残差连接
        x = self.fc1(features)  # 对每个点的特征向量进行线性变换，将其从d_points维转换为d_model维
        q, k, v = self.w_qs(x), index_points(self.w_ks(x), knn_idx), index_points(self.w_vs(x), knn_idx)  # 分别生成查询向量(q)，邻居点向量(k)，邻居点特征向量(v)

        pos_enc = self.fc_delta(xyz[:, :, None] - knn_xyz)  # 学习每个邻居点与中心点之间的位置关系（相对距离），用于加权求和
                              # [batch_size, n_points, n_neighbors, 3] - [batch_size, n_points, n_neighbors, 3] = [batch_size, n_points, n_neighbors, 3]，即为中心点到每个邻居点的相对距离
                              # [batch_size, n_points, n_neighbors, 3] 经过self.fc_delta之后，变为 [batch_size, n_points, n_neighbors, d_model]
        
        attn = self.fc_gamma(q[:, :, None] - k + pos_enc)  # 计算每个邻居点的权重，加上相对位置编码
        attn = F.softmax(attn / np.sqrt(k.size(-1)), dim=-2)  # softmax归一化，保证每个中心点在所有邻居点的权重之和为1

        res = torch.einsum('bmnf,bmnf->bmf', attn, v + pos_enc)  # 加权求和，算上中心点到邻居点的相对位置编码
                                       # attn的shape为 [batch_size, n_points, n_neighbors, d_model]
                                       # v的shape为 [batch_size, n_points, n_neighbors, d_model]
                                       # 输出res的shape为 [batch_size, n_points, d_model]
        res = self.fc2(res) + pre  # 线性变换，将加权求和的结果从d_model维转换为d_points维，并加上残差连接
        return res, attn

        通过 MLP 与位置编码引入了空间信息，并利用基于注意力机制的 Transformer 模型实现了关系学习任务，能够自适应地利用输入的局部点云信息学习全局特征表示。