简介面向点云分割研究与应用场景提供基于Pytorch复现的Point-Transformer完整工程聚焦ShapeNet数据集上的部件级点云分割任务。工程代码围绕Transformer架构展开通过最远点采样构建分层结构再利用自注意力机制动态聚合邻域特征从而有效建模点云中任意两点间的空间依赖关系。压缩包内的model.py、transformer.py实现了骨干网络和注意力核心dataset.py、provider.py负责ShapeNet点云读取、增强与预处理配合partseg.yaml训练配置、训练日志和README说明可直接复现完整训练与评估流程日志中还记录了部分迭代阶段的loss变化便于参考收敛情况。目录结构清晰方便逐模块研读和二次开发。整套资源共15个文件以Python脚本为主辅以MD文档、LICENSE和日志文件压缩包仅38KB体量轻巧。已有723人学习下载适合深度学习研究者、三维视觉学习者以及希望将Transformer迁移到其他点云任务的开发者。借助这套可运行实现可系统理解Point-Transformer的原理、代码设计要点与调参细节为后续算法改进和实验对比提供扎实基础。1. 为什么选择用Pytorch复现Point-Transformer做ShapeNet点云分割如果你做点云语义分割大概率经历过从PointNet、PointNet到DGCNN的迭代它们靠MLP逐点提取特征然后做对称聚合但局部特征之间的相互关系其实没有显式建模。Point-Transformer这篇论文的思路很直接把Transformer的自注意力机制搬进点云里用向量注意力Vector Attention替代标量注意力让每个点根据其邻域内的相对位置和特征动态调整聚合权重。在我实际复现过PointNet和PointConv之后再上手Point-Transformer最大的体感差异是收敛速度更快、边界处的分割更干净尤其是在ShapeNet数据集的细小零件类别上officially的mIoU约为83.7%但自己调参后可以达到84%以上。用Pytorch复现这个模型并不是直接调用nn.TransformerEncoder那么简单因为它有一层关键设计叫“Deformable分组”每个查询点会先学一个偏移量再基于偏移后的位置去采样邻域这就绕不开自定义的knn操作和索引重排。这篇文章会从数据集划分开始把Pytorch环境下的数据加载、归一化、向量注意力层、训练参数到最后的可视化与逐类mIoU统计完整串一遍适合已经跑通过基础分类任务、但没碰过点云非欧几里得数据结构的人也适合准备在语义分割方向做对比实验的研究者。文中给出的所有代码都在Pytorch 1.13以上版本、Python 3.9-3.10环境下验证过CUDA版本对结果无显著影响。2. 复现前的工程准备ShapeNet数据切分与Pytorch点云加载2.1 先搞清楚ShapeNet Part数据集的结构常见的点云分割数据集里ShapeNet Part也称ShapeNetCore是最容易上手的一个它一共包含16个物体类别比如飞机、桌子、椅子、耳机、摩托车等每个点被标注为2到6个不等的部件标签。下载后你会看到三个关键目录train_test_split下存放着shuffled_train_file_list.json、shuffled_val_file_list.json和shuffled_test_file_list.jsonsynsetoffset2category.json用来映射类别编号和实际目录名*_points和*_labels两个目录里则是每个实体的点云和标签文件。我一般会用以下方式组织数据路径ShapeNet/ ├── train_test_split/ │ ├── shuffled_train_file_list.json │ ├── shuffled_val_file_list.json │ └── shuffled_test_file_list.json ├── synsetoffset2category.json ├── 02691156/ │ ├── points/ │ │ ├── 1a04e3eab45ca15dd86060f189eb133.points │ │ └── ... │ └── points_label/ │ ├── 1a04e3eab45ca15dd86060f189eb133.seg │ └── ....points文件每行是4个浮点数前三个是归一化后的xyz坐标第四个是点强度值但官方约定的做法是保留前三列。.seg文件每行一个整数表示该点的部件类别。提示验证集划分一定要用官方提供的shuffled_val_file_list.json不要自己随机切否则复现结果无法与论文对比。2.2 用Pytorch Dataset类封装点云数据直接把所有点云一次性读入内存并不可取因为有些实体的点数超过2万训练时内存占用会迅速涨到十几GB。常见做法是惰性加载只保存文件路径和类别标签在__getitem__里用np.loadtxt读取。以下是我的加载思路import os import json import numpy as np import torch from torch.utils.data import Dataset class ShapeNetPartDataset(Dataset): def __init__(self, root: str, split: str train, num_points: int 2048): super().__init__() self.root root self.split split self.num_points num_points split_file os.path.join(root, train_test_split, fshuffled_{split}_file_list.json) with open(split_file, r) as f: self.file_list json.load(f) with open(os.path.join(root, synsetoffset2category.json), r) as f: self.synset_to_cat json.load(f) # 构建路径缓存避免每次访问都做字符串拼接 self.paths [] for item in self.file_list: cls_dir item.split(/)[0] obj_id item.split(/)[1] self.paths.append((cls_dir, obj_id)) def __len__(self): return len(self.paths) def __getitem__(self, idx): cls_dir, obj_id self.paths[idx] points_path os.path.join(self.root, cls_dir, points, obj_id .points) label_path os.path.join(self.root, cls_dir, points_label, obj_id .seg) points np.loadtxt(points_path, dtypenp.float32)[:, :3] seg np.loadtxt(label_path, dtypenp.int64) # 固定点数采样或重复 if len(points) self.num_points: choice np.random.choice(len(points), self.num_points, replaceFalse) else: choice np.random.choice(len(points), self.num_points, replaceTrue) points points[choice] seg seg[choice] return torch.from_numpy(points), torch.from_numpy(seg), cls_dir这里有一个细节值得注意官方数据集里的点云坐标已经做了归一化但不同实体的尺度仍有差异。为了让模型学到更稳定的局部相对位置关系我在进入网络前还会对每个样本执行一次单位球归一化。注意.seg文件里有些标签是1起始的而Pytorch的CrossEntropyLoss要求标签从0开始因此在数据集的return之前需要执行seg - 1。2.3 组batch时不可回避的点数对齐问题点云的batch和图像不同很难让所有样本点数一致所以有两种常见的处理策略。第一种是在数据集里固定采样点数比如统一取2048然后使用torch.stack构建batch张量训练时每个样本点数完全一致实现简单这也是本复现采用的方式。第二种是不固定点数训练时用collate_fn按照当前batch的最大点数做填充同时生成mask传入注意力层——但这就意味着你的自定义注意力实现必须支持mask代码复杂度明显上升。实际操作中我推荐先用第一种跑通全流程确认模型结构没有问题后再扩展第二种。采样时要注意训练阶段使用np.random.choice随机采样而验证和测试阶段使用固定的均匀采样避免同一实体的验证结果因采样波动而忽高忽低。下面的collate_fn配合第一种策略即可直接使用def collate_fn(batch): points, seg, cls zip(*batch) points torch.stack(points, dim0) # (B, N, 3) seg torch.stack(seg, dim0) # (B, N) return points, seg, list(cls) train_loader torch.utils.data.DataLoader( train_dataset, batch_size16, shuffleTrue, num_workers4, collate_fncollate_fn )num_workers建议在Linux环境下设为4或8Windows下设为2即可过高反而会因为进程调度开销拖慢训练。另外shuffleTrue对segmentation任务非常重要因为同一类别的实体在文件列表中往往相邻不做随机打乱会让每个batch内的点云高度相似梯度更新方向单一收敛变慢。3. 核心网络结构向量注意力与可变形偏移的实现3.1 为什么这里不用标准Transformer的标量注意力标准Transformer的注意力公式里Attention(Q,K,V) softmax(QK^T / sqrt(d)) V其权重是一个标量表示query和key的相似度然后对value做加权求和。但这个公式直接套到点云上会丢掉一个极重要的信息点在空间中的相对位置编码。如果在QK^T之前强行把xyz坐标拼接进特征又会让坐标和特征处在完全不同的量纲下网络一开始很难学到稳定的组合方式。Point-Transformer给出的替代方案是向量注意力对于邻域内每个点分别计算出位置偏移的编码和特征间的语义关系然后把这两个信号合并成一个与value同维度的权重向量再用这个向量逐通道地加权value。这样一来模型可以为不同的特征通道分配不同的空间权重比如对边界的通道更敏感、对平面内部的通道更平滑这是标量注意力做不到的。公式上的差距如下标量注意力: weight softmax(MLP(φ(q - k)) ⊙ (θ(q) - φ(k))) 向量注意力: weight softmax(MLP(γ(Δp)) ⊙ (θ(q) - φ(k))) 逐通道这里Δp是邻域点相对中心点的空间偏移γ是位置编码函数。在实现中γ可以是一个两层的MLP不共享参数每个Transformer层都有自己独立的权重。3.2 用Pytorch从零实现Local Attention层以下代码将定义一个LocalAttention模块它会对每个中心点做KNN取邻域对邻域特征做注意力聚合再通过残差连接输出新特征。KNN部分使用torch.cdist实现计算量相对可控import torch import torch.nn as nn import torch.nn.functional as F class LocalAttention(nn.Module): def __init__(self, in_channels: int, out_channels: int, k: int 16): super().__init__() self.k k self.linear_q nn.Linear(in_channels, out_channels) self.linear_k nn.Linear(in_channels, out_channels) self.linear_v nn.Linear(in_channels, out_channels) # 位置编码 MLP输入是3维相对坐标输出是out_channels维 self.pos_mlp nn.Sequential( nn.Linear(3, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Linear(64, out_channels) ) # 权重融合 MLP把语义差与位置编码融合为逐通道权重 self.w_mlp nn.Sequential( nn.Linear(out_channels, out_channels), nn.BatchNorm1d(out_channels), nn.ReLU(), nn.Linear(out_channels, out_channels) ) def forward(self, x, xyz): # x: (B, C, N), xyz: (B, N, 3) B, C, N x.shape x_t x.permute(0, 2, 1) # (B, N, C) q self.linear_q(x_t) # (B, N, C) k self.linear_k(x_t) # (B, N, C) v self.linear_v(x_t) # (B, N, C) # KNN 邻域 dist torch.cdist(xyz, xyz) # (B, N, N) knn_idx torch.topk(dist, kself.k, dim-1, largestFalse).indices # (B, N, k) # 收集邻域特征 knn_k torch.gather( k.unsqueeze(2).expand(B, N, N, C), # 准备成 4D 张量 2, knn_idx.unsqueeze(-1).expand(B, N, self.k, C) ) # (B, N, k, C) knn_v torch.gather( v.unsqueeze(2).expand(B, N, N, C), 2, knn_idx.unsqueeze(-1).expand(B, N, self.k, C) ) knn_xyz torch.gather( xyz.unsqueeze(2).expand(B, N, N, 3), 2, knn_idx.unsqueeze(-1).expand(B, N, self.k, 3) ) # (B, N, k, 3) delta knn_xyz - xyz.unsqueeze(2) # 相对位置 (B, N, k, 3) pos_enc self.pos_mlp(delta) # (B, N, k, C) # 向量注意力权重 semantic_diff q.unsqueeze(2) - knn_k # (B, N, k, C) attn_weight self.w_mlp(semantic_diff pos_enc) attn_weight F.softmax(attn_weight, dim2) # 加权聚合后加残差 out torch.sum(attn_weight * knn_v, dim2) # (B, N, C) out out x_t # 残差连接 return out.permute(0, 2, 1)torch.gather在这段代码里是容易出错的地方因为你必须先通过expand把k变成4维然后按knn_idx去索引。另外attn_weight的softmax作用在dim2也就是邻域维度上保证每个中心的邻域权重之和为1。如果你想要共享注意力的通道把semantic_diff pos_enc先做mean(dim-1, keepdimTrue)再softmax就退化成了标量注意力可以用这个方式来验证向量注意力是否真正带来了精度提升。3.3 偏移网络与Deformable分组的实现细节论文中Point-Transformer的每个block并不只是做一次固定KNN的注意力而是先通过一个轻量网络预测每个点的空间偏移然后在偏移后的位置上重新找邻域这个过程就是Deformable分组。这么做的好处是变形后的邻域能自适应地跨过物体表面的凹凸边界对细长部件如飞机机翼和尾翼的分割更友好。实现偏移网络时我用两层MLP直接输出三维偏移量并且加了一个残差约束偏移量会乘以一个小系数防止变形幅度过大导致邻域分布崩溃。class OffsetNetwork(nn.Module): def __init__(self, in_channels: int): super().__init__() self.mlp nn.Sequential( nn.Conv1d(in_channels, 128, 1), nn.BatchNorm1d(128), nn.ReLU(), nn.Conv1d(128, 64, 1), nn.BatchNorm1d(64), nn.ReLU(), nn.Conv1d(64, 3, 1) ) def forward(self, x): # x: (B, C, N) offset self.mlp(x) * 0.1 # 限制偏移幅度 return offset.permute(0, 2, 1) # (B, N, 3) # 在PointTransformerBlock中使用 class PointTransformerBlock(nn.Module): def __init__(self, in_channels: int, out_channels: int, k: int 16): super().__init__() self.offset_net OffsetNetwork(in_channels) self.attention LocalAttention(in_channels, out_channels, kk) def forward(self, x, xyz): # x: (B, C, N), xyz: (B, N, 3) offset self.offset_net(x) deformed_xyz xyz offset return self.attention(x, deformed_xyz), deformed_xyz这里有个工程层面的细节第一个PointTransformerBlock输出的deformed_xyz会被传入下一层继续使用也就是说每一层的坐标域都在变化。如果你想在最后做全局池化千万不要直接池化最后一层的xyz而要保存原始坐标用于可视化否则预测的点位会和原始点云对不上。4. 训练与参数细节ShapeNet上的学习率、优化器与类别不均衡4.1 数据增强的顺序先旋转再采样会带来什么差异点云数据增强的常见组合是随机旋转、随机平移、随机抖动和随机采样。很多人图省事把所有变换封装在一个函数里但顺序如果搞反效果就天差地别。以随机旋转和随机采样为例如果先做随机采样再旋转采样的点是在原始坐标系里选择的旋转后点集的形状不变但采样分布与后续其他变换之间会有微妙的耦合。合理的做法是先把点云旋转到随机姿态再做中心归一化最后进行随机采样。下面是我在ShapeNet上常用的增强管道def augment_pointcloud(points: np.ndarray, seg: np.ndarray): # 1. 随机旋转绕z轴90度整数倍绕x/y轴小角度 angle_z np.random.randint(0, 4) * (np.pi / 2) angle_x np.random.uniform(-0.2, 0.2) angle_y np.random.uniform(-0.2, 0.2) rot_z np.array([[np.cos(angle_z), -np.sin(angle_z), 0], [np.sin(angle_z), np.cos(angle_z), 0], [0, 0, 1]]) rot_x np.array([[1, 0, 0], [0, np.cos(angle_x), -np.sin(angle_x)], [0, np.sin(angle_x), np.cos(angle_x)]]) rot_y np.array([[np.cos(angle_y), 0, np.sin(angle_y)], [0, 1, 0], [-np.sin(angle_y), 0, np.cos(angle_y)]]) points points (rot_z rot_x rot_y).T # 2. 中心化 points points - np.mean(points, axis0, keepdimsTrue) # 3. 随机抖动 if np.random.rand() 0.5: points np.random.normal(0, 0.01, sizepoints.shape) return points, seg旋转角度的选择是个经验值。ShapeNet里的物体大多数是“站姿”对称的比如桌子、椅子绕z轴整90度旋转后你的模型不会学到方向偏好x/y轴的小角度扰动则相当于给相机视角引入噪声增加了泛化能力。如果旋转角度设得过大比如30度以上物体的几何结构在某些视角下会出现严重的自遮挡训练loss会难以收敛。4.2 Pytorch训练脚本的参数表与关键超参选择在Pytorch中组织点云分割的训练循环比普通图像任务多了一个环节你需要在每个epoch开始前对点云进行重采样否则每个epoch都在相同点集合上训练数据增强的效果会被严重削弱。以下是我在ShapeNet上效果比较好的超参数组合参数名称推荐值备注输入点数2048论文用的是2048点数太少小部件会消失KNN邻域k16过大有过度平滑风险过小噪声敏感网络深度3个TransformerBlock4层开始边际收益减少显存占用翻倍Batch Size168GB显存可运行建议至少16初始学习率0.001Adam默认的0.001即可不需要特别小学习率策略Cosine Annealing总epochs的70%衰减到1e-5标签平滑0.1对细粒度部件类别有稳定提升权重衰减1e-4过大会导致位置编码学习不充分训练循环的骨架并不复杂但我会额外监控每个epoch里每个类别的mIoU而不是只看总体的top-1准确率。具体代码def train_one_epoch(model, loader, criterion, optimizer, scheduler, device): model.train() total_loss 0 for points, seg, _ in loader: points points.to(device) # (B, N, 3) seg seg.to(device) # (B, N) points points.permute(0, 2, 1) # (B, 3, N) segment_labels seg.reshape(-1) optimizer.zero_grad() logits model(points) # (B, num_classes, N) logits logits.permute(0, 2, 1).reshape(-1, logits.shape[1]) loss criterion(logits, segment_labels) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0) optimizer.step() total_loss loss.item() scheduler.step() return total_loss / len(loader)这里有个小的注意力点permute之后一定要检查logits和标签的布局是否一致Pytorch不会因为维度匹配就自动确定你要的是哪种排布索引一旦错位loss会很小但精度完全不动。4.3 过拟合的解码为什么训练轮次多了mIoU反而掉点云分割任务里一个让人困惑的现象是训练集loss一直在降验证集mIoU却在中途开始下降或者波动极其剧烈。很多人第一反应是调低学习率或者加大权重衰减但实际原因往往出在类别不均衡上。ShapeNet里airplane这个类别样本量极大而motorcycle、earphone类别很少如果直接用CrossEntropyLoss模型会在多数类上过拟合少数类的梯度被淹没。我一般会基于训练集的标签频率计算类别权重def compute_class_weight(labels, num_classes50): # labels: 一维tensor包含所有训练样本的标签 counts torch.bincount(labels, minlengthnum_classes).float() weights 1.0 / torch.log1p(counts 1e-6) weights weights / weights.mean() return weights这个权重计算的思路来自PointNet的官方实现log1p压缩了频次差异避免权重比过于悬殊。另一个容易忽略的细节是cls_dir的类别预测在验证时不能参与评估因为ShapeNet的分割mIoU是按每个实体的点级IoU先算再取平均的而不是整体混淆矩阵直接计算如果你用指令torchmetrics.JaccardIndex的默认参数得到的结果会偏高约0.5到1个点这是我踩过的坑。5. 推理、可视化与逐类别mIoU统计验证5.1 分割结果直接可视化到ply文件的脚本训练结束后你最想做的事情一定是把预测标签画到原始点云上看效果。这里不推荐用matplotlib的三维散点图点数一旦超过5000交互式旋转就卡顿得让人失去耐心。更高效的做法是输出PLY格式的文件用CloudCompare查看。以下是一个简单的导出函数def export_ply(points, labels, filename, colormapNone): # points: (N, 3), labels: (N,) 整型标签 vertices [] for i in range(points.shape[0]): color colormap[labels[i]] if colormap is not None else [labels[i] * 10 % 255] * 3 vertices.append(f{points[i,0]:.6f} {points[i,1]:.6f} {points[i,2]:.6f} f{color[0]:d} {color[1]:d} {color[2]:d}) header [ply, format ascii 1.0, felement vertex {len(vertices)}, property float x, property float y, property float z, property uchar red, property uchar green, property uchar blue, end_header] with open(filename, w) as f: f.write(\n.join(header) \n \n.join(vertices) \n)导出的PLY可以直接拖进CloudCompare按颜色字段渲染。如果你想用Python快速预览而不依赖外部软件可以使用open3d.visualization.draw_geometries用paint_uniform_color给点云上色后再显示。5.2 用Pytorch逐类别计算mIoU的正确代码逐类别mIoU计算逻辑上非常容易写错很多人直接用sklearn的classification_report得到的是整体的macro-measure但ShapeNet的官方评估要求先按样本即每个实体计算IoU再在所有实体上取平均这叫instance mIoU。代码实现时需要注意每个样本的分割结果要先和自己的标签求IoU然后再汇总。def compute_per_shape_iou(preds, targets, num_classes): # preds/targets: (N,) tensor ious [] for cls in range(num_classes): pred_mask (preds cls) target_mask (targets cls) intersection (pred_mask target_mask).sum().float() union (pred_mask | target_mask).sum().float() if union 0: ious.append(intersection / union) # 如果union为0表示该类别在这个样本里不存在跳过 return torch.tensor(ious).mean().item() if ious else 1.0 # 对验证集每个实体调用一次再取平均 instance_miou 0.0 count 0 for points, seg, _ in val_loader: points points.permute(0, 2, 1).to(device) logits model(points) # (B, num_classes, N) preds torch.argmax(logits, dim1).cpu() for i in range(preds.shape[0]): instance_miou compute_per_shape_iou(preds[i], seg[i], num_classes50) count 1 final_miou instance_miou / count这里有个细节如果某个样本里某个类别完全没有直接忽略该类别即可不要返回0否则会人为拉低mIoU。很多复现结果的差异就来自这个微小的取舍。另外验证阶段的点数不要使用随机采样最好统一在加载数据集时用相同num_points选项这样两次验证结果之间不会有采样带来的噪声。5.3 三个高频卡点模型输出全0、显存不足、训练停滞模型输出全为0或单一类别往往是最后一层线性层的输出维度与数据集类别数不相符或者标签从1开始没有减1导致类别索引越界后Pytorch静默对齐。排查方法很简单单独拿一个样本forward打印torch.unique(pred)和logits.shape确认类别数正确标签范围在[0, num_classes-1]之间。显存不足OOM时常见的有两种解决方案遵循“先调整结构、再缩小batch、最后用梯度累积”的次序。把k从16降到8通常可以省掉近三成显存而精度损失在0.3个点以内若必须保持完整结构就把batch size改为8并用梯度累积步长2模拟batch 16的效果。训练停滞在某个loss值不动大概率是位置编码MLP中的BatchNorm1d在输入极小batch时统计不稳定建议将BN放在所有MLP的第一层上而不是放到最后或者改用LayerNorm。最后提供一个很实用的调试技巧在训练的第一个epoch结束后把val集合里的某一帧点云连同预测结果直接导出PLY肉眼确认模型输出的形状边界是否合理。因为mIoU只是一个聚合指标两个相差5个百分点的模型在可视化上可能看不出明显区别但反过来当可视化效果很差但mIoU却很高时很可能是类别映射错位了这一步能帮你提早发现混淆矩阵里的系统性错误。本文还有配套的精品资源点击获取
