PyTorch Geometric 点云分类实战:从 DGCNN 到多卡训练

PyTorch Geometric 点云分类实战:从 DGCNN 到多卡训练 PyTorch Geometric 点云分类实战从 DGCNN 到多卡训练【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric给一个 1024 个点的三维点云分类最直觉的做法是塞进 PointNet 式的全局 MLP——但纯全连接对局部几何不敏感细节相近的类别比如 ModelNet40 里的 chair 和 stool很容易混。一个常用的折中在点集上动态建 K 近邻图让卷积逐层长出局部结构。PyTorch GeometricPyG把这个思路做成了开箱即用的算子DynamicEdgeConv和PointTransformerConv两个类就能分别搭出 DGCNN 和 Point Transformer 两条完整分类管线示例代码都在examples/下几百行可跑通。这套库的定位是 PyTorch 上的图神经网络工具箱对点云任务有三点直接可用的能力torch_geometric.datasets内置 ModelNet、MedShapeNet 等点云数据集下载、切分、转换一行搞定torch_geometric.nn提供 FPS 采样、kNN 建图、动态边卷积、点 Transformer 卷积等算子稀疏张量和 scatter 聚合是原生实现批内点数不齐也不会炸三步搭好数据加载安装上只要两件事装 PyG 本体再装 pyg-libFPS、kNN 这些 C 算子的宿主pip install torch_geometric pip install pyg-lib # DynamicEdgeConv、fps 都依赖它0.6.0更细的依赖矩阵CUDA 版本、可选扩展看官方文档 docs/source/install/ 即可不用逐行抄。数据侧的套路是预变换 在线变换两段式pre_transform, transform T.NormalizeScale(), T.SamplePoints(1024) # NormalizeScale 把点云中心化并缩放到单位球——跨样本尺度一致卷积才稳定 # SamplePoints(1024) 在线采样到固定点数DynamicEdgeConv 建图才不会 OOM train_dataset ModelNet(./data/modelnet10, 10, True, transform, pre_transform) test_dataset ModelNet(./data/modelnet10, 10, False, transform, pre_transform)完整的数据集加载逻辑含 MedShapeNet 按类别 7:3 切分见 examples/dgcnn_classification.py#L42-L76。DGCNN图不是给好的是每次 forward 现算的DGCNN 的核心类是DynamicEdgeConv源码在 torch_geometric/nn/conv/edge_conv.py每层前向传播时它都会基于当前坐标或特征重新做一次 kNN 建图然后对每条边做点 i 的特征 ⊕ (x_i − x_j) 的差向量再 MLP。差向量是关键——它让消息只编码相对几何天然平移不变。示例里的完整网络只有两层动态卷积见 examples/dgcnn_classification.py#L91-L108class Net(torch.nn.Module): def __init__(self, out_channels, k20, aggrmax): super().__init__() self.conv1 DynamicEdgeConv(MLP([2*3, 64, 64, 64]), k, aggr) # k20每个点聚合 20 个近邻 self.conv2 DynamicEdgeConv(MLP([2*64, 128]), k, aggr) self.lin1 Linear(128 64, 1024) self.mlp MLP([1024, 512, 256, out_channels], dropout0.5, normNone) def forward(self, data): pos, batch data.pos, data.batch x1 self.conv1(pos, batch) # 图由 3D 坐标现算 x2 self.conv2(x1, batch) # 图由上一层特征现算 out self.lin1(torch.cat([x1, x2], dim1)) # 两层特征拼接避免浅层信息丢失 out global_max_pool(out, batch) # 变长点集 max 池化成单向量 return F.log_softmax(self.mlp(out), dim1)注意conv2吃的是特征x1而不是原始坐标——第二层图已经建在特征空间里这是动态的含义。k 值从 20 起步点云更稀疏或想省显存时调到 12~15细节类任务MedShapeNet调大到 25 通常有小幅收益但 kNN 的计算量随 k 线性涨。Point Transformer注意力 逐级降采样另一条路线是PointTransformerConvtorch_geometric/nn/conv/point_transformer_conv.py注意力权重由特征差 相对位置嵌入共同决定位置信息显式进了 attention 的计算。示例 examples/point_transformer_classification.py 的骨架是Transformer 块 TransitionDown 降采样交替堆叠dim_model[32, 64, 128, 256, 512]五级。降采样部分是整篇最值得抄的片段它把 FPS 和 kNN 组合成了标准套路examples/point_transformer_classification.py#L64-L84id_clusters fps(pos, ratio0.25, batchbatch) # FPS最远点采样每步挑离已选点最远的那个0.25 即保留 1/4 点覆盖均匀 sub_batch batch[id_clusters] id_k_neighbor knn(pos, pos[id_clusters], k16, batch_xbatch, batch_ysub_batch) # kNN给每个采样中心找 16 个原始点聚合范围不丢 x_out scatter(x[id_k_neighbor[1]], id_k_neighbor[0], dim0, reducemax) # 邻域取 max 作为中心特征点数逐级 1024→256→64→16 return x_out, pos[id_clusters], sub_batch这条路径比 DGCNN 多一套注意力开销但下采样后深层的计算量反而更低两个示例都在 ModelNet10 上跑 201 个 epoch、batch size 32、StepLR(step_size20, gamma0.5)直接在同一数据集上对比公平。训练循环与数字锚点训练循环没有花活就是一个 nll_loss 的循环examples/dgcnn_classification.py#L117-L147python examples/dgcnn_classification.py --dataset modelnet10 --batch_size 32几个可以拿去对齐的数字每样本固定 1024 点、每层 k20、Adam lr0.001、每 20 个 epoch 学习率减半、共 201 个 epoch。单卡上 DGCNN 的瓶颈几乎全在每层的 kNN 建图C 扩展里跑CPU/GPU 均可Point Transformer 则在 attention 上多花一份——显存紧张时先把 batch 从 32 降到 16比砍 k 更不伤精度。MedShapeNet 这类按类别切分的长尾数据示例里用random.seed(42)保证 7:3 划分可复现自己写切分时别漏这一步。往多卡上扩采样并行与分片单卡吃完 1024 点数据集后真正的扩展发生在超大图上PyG 把邻居采样做成 RPC 服务每张卡只持有本地分片采样时跨卡取远程邻居示意如下图来自 docs/source/notes/ 的分布式章节入口是 examples/multi_gpu/distributed_sampling.py用torchrun起多进程后DistNeighborLoader会按分片表自动把本地点和远程点拆好模型侧不用感知——这套机制对点云场景意味着百万级点云可以按空间分片喂给多卡而不必先塞进单卡显存。选型建议与延伸类别靠局部形状区分椅子 vs 桌子DGCNN 够用k 值是第一调参旋钮需要长程依赖或全局结构细粒度、遮挡多Point Transformer接受注意力开销想快速复现对比两个示例共用 ModelNet10/1024 点/201 epoch 设定直接并排跑延伸阅读examples/dgcnn_segmentation.py 同构网络做点级分割global_pool 换成逐点输出即可benchmark/points/ PointNet、EdgeCNN、SplineCNN 在 ModelNet10 上的统一评测脚本examples/multi_gpu/ 分布式采样与模型并行的完整可跑示例有类似点云上叠注意力或多卡切分点云的场景欢迎聊聊踩过的坑。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考