PointNet与PointNet++实战对比：从零搭建三维点云分割模型（附PyTorch代码）

PointNet与PointNet++实战对比：从零搭建三维点云分割模型（附PyTorch代码）

如果你正在为三维点云数据头疼，想找一个能直接上手、效果又好的分割模型，那么PointNet和PointNet++这两个名字你一定绕不开。几年前我第一次接触自动驾驶点云分割项目时，面对一堆无序的XYZ坐标，传统的图像处理方法完全失效，正是这两个模型打开了新世界的大门。它们不像体素方法那样笨重，也不像多视图投影那样丢失信息，而是直接“硬刚”原始点云，这种思路在当时堪称革命。但理论归理论，真正要把它们用起来，尤其是用PyTorch从零搭建，你会发现两者在代码结构、训练技巧和实际表现上差异巨大。这篇文章，我就结合自己踩过的坑和调参的经验，带你深入代码层面，看看PointNet和PointNet++到底该怎么选、怎么用。

1. 环境搭建与数据准备：万事开头难

在动手写模型之前，一个干净、可复现的环境是高效开发的基础。我强烈建议使用Anaconda来管理Python环境，它能很好地解决依赖冲突问题。

# 创建并激活一个名为pointnet的虚拟环境
conda create -n pointnet python=3.8
conda activate pointnet

# 安装核心依赖
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install numpy scipy matplotlib tqdm
pip install open3d  # 用于点云可视化，非常有用

这里我固定了PyTorch的版本，因为不同版本间一些API可能有细微变动，为了代码的稳定性，锁定版本是个好习惯。CUDA版本请根据你自己的显卡驱动进行选择。

接下来是数据。对于点云分割，斯坦福大学的S3DIS（Stanford 3D Indoor Spaces）数据集是一个经典的室内场景分割基准。它包含了6个大型室内区域的3D扫描，每个点都被标注为13个类别（如天花板、地板、桌子等）。

注意：直接处理原始的S3DIS数据（.txt格式）比较低效。一个常见的预处理步骤是将每个房间的点云切割成重叠的1m x 1m的块（block），这样既能控制输入规模，又能通过数据增强增加样本多样性。

下面是一个简化的数据加载器示例，展示了如何读取预处理后的.npy数据块：

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

class S3DISDataset(Dataset):
    def __init__(self, data_root, split='train', block_size=1.0, num_point=4096):
        self.data_root = data_root
        self.split = split
        self.block_size = block_size
        self.num_point = num_point
        # 假设数据已预处理为.npy文件，存储在data_root/split目录下
        self.file_list = self._load_file_list()

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, idx):
        # 加载一个数据块：包含点坐标和标签
        data = np.load(self.file_list[idx])  # shape: (N, 4) -> x, y, z, label
        points, labels = data[:, :3], data[:, 3].astype(np.int64)

        # 数据增强（仅在训练时）
        if self.split == 'train':
            points = self._augment_point_cloud(points)

        # 采样固定数量的点（如果N > num_point，则随机下采样；如果N < num_point，则重复采样）
        choice = np.random.choice(len(points), self.num_point, replace=len(points) < self.num_point)
        points = points[choice, :]
        labels = labels[choice]

        # 归一化：将块中心移到原点
        points -= np.mean(points, axis=0, keepdims=True)

        # 转换为Tensor
        points = torch.from_numpy(points).float()
        labels = torch.from_numpy(labels).long()
        return points, labels

    def _augment_point_cloud(self, points):
        """简单的数据增强：随机旋转和抖动"""
        # 绕Z轴随机旋转
        theta = np.random.uniform(0, 2*np.pi)
        rotation_matrix = np.array([
            [np.cos(theta), -np.sin(theta), 0],
            [np.sin(theta), np.cos(theta), 0],
            [0, 0, 1]
        ])
        points = np.dot(points, rotation_matrix)
        # 添加微小抖动
        jitter = np.random.normal(0, 0.02, size=points.shape)
        points += jitter
        return points

这个Dataset类封装了数据读取、采样、归一化和增强的核心逻辑。使用DataLoader进行批量加载时，由于每个样本的点数已经通过采样固定为num_point，所以可以直接进行堆叠（stack），无需使用pad_sequence。

2. PointNet核心实现：对称函数的艺术

PointNet的核心思想非常优雅：它通过一个共享权重的多层感知机（MLP）独立处理每个点，然后使用一个对称函数（如最大池化）来聚合所有点的信息，从而解决点云的无序性问题。此外，它引入了T-Net来学习一个变换矩阵，对齐输入点云或特征，提升模型的旋转不变性。

我们先来看最关键的模块——T-Net。它的目标是学习一个变换矩阵（对于输入是3x3，对于特征可以是64x64），使网络对空间变换更鲁棒。

import torch.nn as nn
import torch.nn.functional as F

class TNet(nn.Module):
    """学习空间或特征变换的模块"""
    def __init__(self, k=3):
        super().__init__()
        self.k = k  # 变换矩阵的维度，输入变换为3，特征变换为64
        self.conv1 = nn.Conv1d(k, 64, 1)
        self.conv2 = nn.Conv1d(64, 128, 1)
        self.conv3 = nn.Conv1d(128, 1024, 1)
        self.fc1 = nn.Linear(1024, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, k*k)

        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(128)
        self.bn3 = nn.BatchNorm1d(1024)
        self.bn4 = nn.BatchNorm1d(512)
        self.bn5 = nn.BatchNorm1d(256)

    def forward(self, x):
        # x shape: (batch_size, k, num_points)
        batch_size = x.size(0)
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        x = torch.max(x, 2, keepdim=True)[0]  # 全局最大池化
        x = x.view(-1, 1024)
        x = F.relu(self.bn4(self.fc1(x)))
        x = F.relu(self.bn5(self.fc2(x)))
        x = self.fc3(x)

        # 将输出reshape为batch_size个k*k的矩阵，并添加单位矩阵作为初始偏置
        iden = torch.eye(self.k, requires_grad=True).repeat(batch_size, 1, 1)
        if x.is_cuda:
            iden = iden.cuda()
        x = x.view(-1, self.k, self.k) + iden
        return x

有了T-Net，我们就可以构建完整的PointNet分割网络了。分割网络需要输出每个点的类别，因此它在提取全局特征后，会将全局特征与每个点的局部特征进行拼接，再通过MLP进行逐点分类。

class PointNetSeg(nn.Module):
    """PointNet语义分割网络"""
    def __init__(self, num_classes=13, num_point=4096):
        super().__init__()
        self.num_point = num_point
        self.input_transform = TNet(k=3)
        self.feature_transform = TNet(k=64)

        # 共享权重的MLP (卷积实现)
        self.conv1 = nn.Conv1d(3, 64, 1)
        self.conv2 = nn.Conv1d(64, 64, 1)
        self.conv3 = nn.Conv1d(64, 64, 1)
        self.conv4 = nn.Conv1d(64, 128, 1)
        self.conv5 = nn.Conv1d(128, 1024, 1)

        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(64)
        self.bn3 = nn.BatchNorm1d(64)
        self.bn4 = nn.BatchNorm1d(128)
        self.bn5 = nn.BatchNorm1d(1024)

        # 分割头：将全局特征与局部特征融合
        self.conv6 = nn.Conv1d(1088, 512, 1)  # 1024+64=1088
        self.conv7 = nn.Conv1d(512, 256, 1)
        self.conv8 = nn.Conv1d(256, 128, 1)
        self.conv9 = nn.Conv1d(128, num_classes, 1)

        self.bn6 = nn.BatchNorm1d(512)
        self.bn7 = nn.BatchNorm1d(256)
        self.bn8 = nn.BatchNorm1d(128)

    def forward(self, x):
        # x shape: (batch_size, 3, num_points)
        batch_size, _, num_points = x.size()

        # 输入变换
        trans_input = self.input_transform(x)
        x = torch.bmm(x.transpose(2, 1), trans_input).transpose(2, 1)

        # 第一层特征提取
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        point_feat = x  # 保存此时的局部特征 (batch_size, 64, num_points)

        # 特征变换
        trans_feat = self.feature_transform(x)
        x = torch.bmm(x.transpose(2, 1), trans_feat).transpose(2, 1)

        # 继续提取更高层特征
        x = F.relu(self.bn3(self.conv3(x)))
        x = F.relu(self.bn4(self.conv4(x)))
        x = F.relu(self.bn5(self.conv5(x)))  # (batch_size, 1024, num_points)

        # 全局特征：最大池化
        global_feat = torch.max(x, 2, keepdim=True)[0]  # (batch_size, 1024, 1)
        global_feat_repeated = global_feat.repeat(1, 1, num_points)  # 复制到每个点

        # 拼接局部特征与全局特征
        x = torch.cat([point_feat, global_feat_repeated], dim=1)  # (batch_size, 1088, num_points)

        # 分割头
        x = F.relu(self.bn6(self.conv6(x)))
        x = F.relu(self.bn7(self.conv7(x)))
        x = F.relu(self.bn8(self.conv8(x)))
        x = self.conv9(x)  # (batch_size, num_classes, num_points)

        x = x.transpose(2, 1).contiguous()  # (batch_size, num_points, num_classes)
        return x, trans_input, trans_feat

在训练时，论文中除了常规的交叉熵损失，还为正则化特征变换矩阵增加了一个辅助损失，希望它接近正交矩阵，以保持特征空间的稳定性。

def feature_transform_regularizer(trans):
    """特征变换矩阵的正则化损失，使其接近正交矩阵"""
    batch_size, k, _ = trans.size()
    I = torch.eye(k, device=trans.device)[None, :, :]
    loss = torch.norm(torch.bmm(trans, trans.transpose(2, 1)) - I, dim=(1, 2))
    return loss.mean()

PointNet的优势与局限：

• 优势：结构简单，参数少，训练快，对点云的无序性和旋转有理论保证的鲁棒性。
• 局限：最大池化聚合全局特征的方式过于粗暴，完全丢失了点的局部邻域结构信息。这在处理复杂场景（如室内场景分割）时，会导致对细节和边界的分割效果不佳，因为模型“看”不到点与点之间的相对位置关系。

3. PointNet++架构解析：分层与局部感知

PointNet++的核心改进在于引入了分层特征学习和局部区域抽象。它模仿了CNN中卷积层和池化层堆叠的思想，通过多次“采样-分组-特征提取”的操作（称为Set Abstraction），逐步扩大感受野，同时保留局部几何信息。

首先，我们需要实现最远点采样（FPS）和球查询（Ball Query），这是构建局部区域的基础。

def farthest_point_sample(xyz, npoint):
    """
    最远点采样 (FPS)
    Args:
        xyz: (B, N, 3) 点云坐标
        npoint: 需要采样的中心点数量
    Returns:
        centroids: (B, npoint) 采样得到的中心点索引
    """
    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
    batch_indices = torch.arange(B, dtype=torch.long).to(device)

    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    return centroids

def query_ball_point(radius, nsample, xyz, new_xyz):
    """
    球查询 (Ball Query)
    Args:
        radius: 搜索半径
        nsample: 每个区域最多采样的点数
        xyz: (B, N, 3) 所有点的坐标
        new_xyz: (B, S, 3) 中心点坐标
    Returns:
        group_idx: (B, S, nsample) 每个中心点邻域内的点索引
    """
    device = xyz.device
    B, N, C = xyz.shape
    _, S, _ = new_xyz.shape
    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
    sqrdists = square_distance(new_xyz, xyz)  # (B, S, N)
    group_idx[sqrdists > radius ** 2] = N  # 将距离大于半径的索引标记为N（无效索引）
    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]  # 取前nsample个最近的
    # 如果某个区域内有效点不足nsample，会有N这个无效索引，后续需要处理
    return group_idx

有了这些基础操作，我们就可以构建PointNet++的核心模块——Set Abstraction (SA) Layer。

class PointNetSetAbstraction(nn.Module):
    """PointNet++的Set Abstraction层"""
    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all=False):
        super().__init__()
        self.npoint = npoint
        self.radius = radius
        self.nsample = nsample
        self.group_all = group_all  # 如果为True，则将所有点作为一个组（用于最后一层）

        # 用于提取局部特征的微型PointNet (MLP)
        self.mlp_convs = nn.ModuleList()
        self.mlp_bns = nn.ModuleList()
        last_channel = in_channel
        for out_channel in mlp:
            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
            last_channel = out_channel

    def forward(self, xyz, points):
        """
        Args:
            xyz: (B, N, 3) 输入点坐标
            points: (B, C, N) 输入点特征，可以为None
        Returns:
            new_xyz: (B, npoint, 3) 采样后的中心点坐标
            new_points: (B, mlp[-1], npoint) 采样后中心点的新特征
        """
        if self.group_all:
            new_xyz, new_points = sample_and_group_all(xyz, points)
        else:
            new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
        # new_points shape: (B, 3+C, npoint, nsample)

        # 使用微型PointNet处理每个局部区域
        new_points = new_points.permute(0, 3, 2, 1)  # (B, nsample, npoint, 3+C) -> (B, 3+C, npoint, nsample)
        for i, conv in enumerate(self.mlp_convs):
            bn = self.mlp_bns[i]
            new_points = F.relu(bn(conv(new_points)))
        # new_points shape: (B, mlp[-1], npoint, nsample)

        # 在nsample维度上进行最大池化，得到每个局部区域的聚合特征
        new_points = torch.max(new_points, 3)[0]  # (B, mlp[-1], npoint)
        return new_xyz, new_points

PointNet++的分割网络是一个编码器-解码器结构，类似于U-Net。编码器通过多个SA层下采样并提取多尺度特征，解码器则通过特征传播（Feature Propagation）层进行上采样，并将编码器对应层的特征进行跳跃连接（Skip Connection），以恢复细节信息。

class PointNet2Seg(nn.Module):
    """PointNet++语义分割网络"""
    def __init__(self, num_classes=13, num_point=4096):
        super().__init__()
        self.sa1 = PointNetSetAbstraction(npoint=1024, radius=0.1, nsample=32, in_channel=3, mlp=[32, 32, 64])
        self.sa2 = PointNetSetAbstraction(npoint=256, radius=0.2, nsample=32, in_channel=64+3, mlp=[64, 64, 128])
        self.sa3 = PointNetSetAbstraction(npoint=64, radius=0.4, nsample=32, in_channel=128+3, mlp=[128, 128, 256])
        self.sa4 = PointNetSetAbstraction(npoint=16, radius=0.8, nsample=32, in_channel=256+3, mlp=[256, 256, 512])

        # 特征传播（上采样）层
        self.fp4 = PointNetFeaturePropagation(in_channel=512+256, mlp=[256, 256])
        self.fp3 = PointNetFeaturePropagation(in_channel=256+128, mlp=[256, 256])
        self.fp2 = PointNetFeaturePropagation(in_channel=256+64, mlp=[256, 128])
        self.fp1 = PointNetFeaturePropagation(in_channel=128+3, mlp=[128, 128, 128])

        # 分割头
        self.conv1 = nn.Conv1d(128, 128, 1)
        self.bn1 = nn.BatchNorm1d(128)
        self.drop1 = nn.Dropout(0.5)
        self.conv2 = nn.Conv1d(128, num_classes, 1)

    def forward(self, xyz):
        # xyz: (B, N, 3)
        B, N, _ = xyz.shape
        l0_points = xyz.transpose(2, 1)  # (B, 3, N)
        l0_xyz = xyz

        # 编码器
        l1_xyz, l1_points = self.sa1(l0_xyz, l0_points)  # (B, 1024, 3), (B, 64, 1024)
        l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)  # (B, 256, 3), (B, 128, 256)
        l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)  # (B, 64, 3), (B, 256, 64)
        l4_xyz, l4_points = self.sa4(l3_xyz, l3_points)  # (B, 16, 3), (B, 512, 16)

        # 解码器（特征传播）
        l3_points = self.fp4(l3_xyz, l4_xyz, l3_points, l4_points)  # (B, 256, 64)
        l2_points = self.fp3(l2_xyz, l3_xyz, l2_points, l3_points)  # (B, 256, 256)
        l1_points = self.fp2(l1_xyz, l2_xyz, l1_points, l2_points)  # (B, 128, 1024)
        l0_points = self.fp1(l0_xyz, l1_xyz, None, l1_points)       # (B, 128, N)

        # 分割头
        x = F.relu(self.bn1(self.conv1(l0_points)))
        x = self.drop1(x)
        x = self.conv2(x)  # (B, num_classes, N)
        x = x.transpose(2, 1).contiguous()  # (B, N, num_classes)
        return x

PointNet++的优势与挑战：

• 优势：通过分层结构捕获了丰富的局部几何信息，在复杂场景的分割任务上精度显著优于PointNet，尤其是对物体边界和细节部分。
• 挑战：计算复杂度高，尤其是FPS和Ball Query操作在点数多时非常耗时。对点云密度的变化敏感，在稀疏区域可能无法有效分组。网络结构更复杂，调参（如半径radius、采样数nsample）需要更多经验。

4. 训练策略与调参实战：从理论到结果

模型搭建好了，但要让它们真正work起来，训练策略和参数调优才是关键。这里我分享一些在S3DIS数据集上训练这两个模型时积累的经验。

优化器与学习率调度：对于这类3D任务，Adam优化器通常是稳妥的起点。学习率初始值可以设为0.001。我习惯使用余弦退火（CosineAnnealingLR）或者带热重启的余弦退火（CosineAnnealingWarmRestarts），它们能让模型在训练后期更精细地收敛。

import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR

model = PointNet2Seg(num_classes=13).cuda()
optimizer = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999), weight_decay=1e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=200, eta_min=1e-5)  # 假设总epoch为200

损失函数：对于分割任务，标准的交叉熵损失（CrossEntropyLoss）是基础。但在点云数据中，类别不平衡问题非常严重（例如，墙和地板点的数量远多于家具）。一个有效的技巧是使用加权交叉熵损失，根据每个类别在训练集中的频率来设置权重。

# 假设我们统计了训练集中13个类别的点数，得到类别权重列表 class_weights
class_weights = torch.tensor([...], dtype=torch.float32).cuda()
criterion = nn.CrossEntropyLoss(weight=class_weights, ignore_index=-1)  # ignore_index用于忽略无效点

数据增强：这是提升模型泛化能力、防止过拟合的利器。除了代码示例中提到的随机旋转和坐标抖动，还可以尝试：

• 随机缩放：以轻微不同的比例缩放点云。
• 随机平移：在三个轴上随机移动点云。
• 随机丢弃点：以一定概率随机丢弃一些点，模拟传感器噪声或遮挡，这对提升鲁棒性很有帮助。

训练过程中的常见问题与解决：

1. 梯度爆炸/消失：如果训练初期loss就变成NaN，很可能是梯度爆炸。可以尝试：降低学习率；使用梯度裁剪（torch.nn.utils.clip_grad_norm_）；检查网络初始化，确保没有过大的权重。
2. 过拟合：如果训练集精度很高但验证集精度停滞不前，就是过拟合。对策包括：增加数据增强的强度；在MLP层后添加Dropout层（如PointNet++分割头所示）；增大权重衰减（weight_decay）；使用早停法（Early Stopping）。
3. PointNet++训练速度慢：FPS和Ball Query是瓶颈。可以尝试：在训练初期使用较小的num_point（如2048），后期再增大；使用CUDA加速的最近邻搜索库（如torch_cluster）；如果显存充足，可以适当增大batch_size以更充分利用GPU。

模型评估与对比：在S3DIS数据集上，我们通常使用**平均交并比（mIoU）**作为主要评估指标。下表展示了一个简单的性能对比框架（数值为示意，实际结果需训练得到）：

模型	整体精度 (OA)	平均交并比 (mIoU)	参数量	推理速度 (pts/s)	适用场景分析
PointNet	~85%	~45%	~3M	快	对简单物体、部件分割效果好，适合计算资源受限、对实时性要求高的场景。
PointNet++	~88%	~55%	~5M	中等	对复杂室内外场景分割更优，能捕捉细节，但计算开销大，需仔细调参。

实际训练中，PointNet可能更快达到一个不错的基线，而PointNet++则需要更长的训练时间和更精细的超参数调整（特别是radius和nsample），但一旦调好，其上限通常更高。

最后，别忘了可视化你的结果。用Open3D这样的库将预测结果和真实标签渲染出来，直观对比，能帮你发现模型在哪里犯了错，是边界模糊、小物体漏检，还是类别混淆，这比只看数字更有助于理解模型的优缺点。