PyTorch Geometric 3D 点云实战：DGCNN 与 Point Transformer 从数据加载到调优

PyTorch Geometric 3D 点云实战DGCNN 与 Point Transformer 从数据加载到调优【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric假设你手头有几十万张 3D 点云样本要按类别识别椅子、桌子、飞机……还要支持逐点分割。手动处理这类数据会非常琐碎点云没有固定的网格结构点数不定卷积核没法直接搬。PyTorch GeometricPyG就是为这类非欧几里得数据准备的图库torch_geometric包里内置了点云数据集、建图变换、DynamicEdgeConv、PointTransformerConv等模型组件以及面向大图的分布式数据加载。本文按先跑通 → 按任务拆 → 调优的顺序带你看清这条链路。能力地图PyG 点云能力的四层结构PyG 在 3D 点云方向的能力可以分成四层从下往上依次是数据层ModelNet、MedShapeNet、ShapeNet、S3DIS等数据集类torch_geometric/datasets/统一以Data对象暴露点坐标存在pos字段变换层T.NormalizeScale、T.SamplePoints(1024)负责归一化和固定点数KNNGraph、RadiusGraph把无序点云变成近邻图见 docs/source/tutorial/point_cloud.rst模型层消息传递基类之上实现DynamicEdgeConvtorch_geometric/nn/conv/edge_conv.py、PointTransformerConvtorch_geometric/nn/conv/point_transformer_conv.py配fps、knn、global_max_pool等池化算子训练/部署层DataLoader自动拼接变长点云并维护batch向量torch_geometric/distributed/提供多机采样加载torch.jit支持 TorchScript 导出点云处理的核心流程如下图所示采样分组 → 邻域聚合 → 降采样循环若干层后得到整云表示。DGCNN、PointNet、Point Transformer 的差异主要体现在邻域聚合这一步用什么算子。最小上手路径ModelNet10 三步跑通第一步安装。除了 PyG 本体外点云模型依赖pyg-lib提供fps、knn、DynamicEdgeConv的 C 后端pip install torch_geometric pyg-lib若同时用torch-scatter、torch-sparse等 CUDA 扩展需按你本地 torch 版本对应的 wheel 安装见 docs/source/get_started/版本不匹配是新手最常见的报错来源。第二步加载数据 写模型 训练下面这个 ≤15 行的最小示例可直接复制运行结构取自 examples/dgcnn_classification.pyimport torch, torch.nn.functional as F import torch_geometric.transforms as T from torch_geometric.datasets import ModelNet from torch_geometric.loader import DataLoader from torch_geometric.nn import MLP, DynamicEdgeConv, global_max_pool pre, tra T.NormalizeScale(), T.SamplePoints(1024) ds ModelNet(./data, 10, True, transformtra, pre_transformpre) dl DataLoader(ds, batch_size32, shuffleTrue) class Net(torch.nn.Module): def __init__(self): super().__init__() self.conv1 DynamicEdgeConv(MLP([6, 64, 64, 64]), k20) self.conv2 DynamicEdgeConv(MLP([128, 128]), k20) self.head MLP([128 64, 256, ds.num_classes], dropout0.5, normNone) def forward(self, data): x1 self.conv1(data.pos, data.batch) x2 self.conv2(x1, data.batch) return F.log_softmax( self.head(global_max_pool(torch.cat([x1, x2], dim1), data.batch)), dim1) model, opt Net(), torch.optim.Adam(Net().parameters(), lr0.001) for data in dl: opt.zero_grad(); loss F.nll_loss(model(data), data.y); loss.backward(); opt.step()跑起来后每个 epoch 约 400 步ModelNet10 训练集 3870 个模型、batch_size32单张消费级 GPU 一个 epoch 数分钟到十几分钟量级上完全可日常迭代。任务拆解先想清楚要解决什么问题点云分类DGCNN 的动态建边思路问题DGCNN 论文的核心是每一层都在特征空间里重新找近邻图的连接关系随训练变化。一句话直觉把它类比成每次卷积都重画一张局部邻居图——第一层用坐标当特征找 KNN之后用学到的特征找 KNN邻居关系跟着语义走。关键 APIDynamicEdgeConv(nn, k, aggrmax)内部自动做 KNN 建边再消息传递你只需给出边权 MLP输入维度是2 * in_channels因为两两拼接conv DynamicEdgeConv(MLP([2 * 3, 64, 64, 64]), k20, aggrmax) x1 conv(data.pos, data.batch) # 第一层用坐标 x2 DynamicEdgeConv(MLP([2 * 64, 128]), 20)(x1, data.batch) out global_max_pool(x2, data.batch) # 整云 max 池化 → 每图一个向量分类头接 2~3 层MLPF.log_softmax即可。效果这套结构在 ModelNet10 上收敛到约 92% top-1论文报告 92.2%是点云分类的基线级选择换成 ModelNet40 会明显更难。分割任务则是同一套卷积去掉全局池化、逐点接分类头可参考 examples/dgcnn_segmentation.pyShapeNet 上按类别训练用 jaccard 指标评估。长程依赖建模与分割Point Transformer 的全局注意力思路问题KNN 局部聚合只看得见 20 个近邻整云级的形状关系比如椅背和椅腿的相对布局要靠堆很深才能传到位。一句话直觉PointTransformerConv把 KNN 子图当作注意力的候选集——每个点只对邻居做 Q/K/V 加权等价于在局部做了一次 softmax 归一化比 max 聚合保留了谁贡献多少的信息。关键 APIPointTransformerConv需要自己用knn_graph提供边再用fps按比例降采样构成编码器式结构完整实现见 examples/point_transformer_classification.pyedge_index knn_graph(pos, k16, batchbatch) # KNN 子图 x PointTransformerConv(dim, dim, pos_nnpos_nn, attn_nnattn_nn)(x, pos, edge_index) id fps(pos, ratio0.5, batchbatch) # FPS 降采样一半点 edge_index knn_graph(pos[id], k16, batchbatch) x x[id]; pos, batch pos[id], batch[id]效果在同样的 ModelNet10 设置下Point Transformer 论文报告约 93% top-1略高于同设置 DGCNN代价是注意力算子比 max 池化更吃显存且该示例要求pyg-lib0.6.0脚本开头会检查WITH_PYG_LIB。分割场景把末层global_mean_pool换成逐点MLP输出即可。进阶与调优分布式、推理、内存三个痛点痛点单机装不下/跑不快大图大点云→ 解法PyG 自带分布式采样框架torch_geometric/distributed/把图按节点分片存到各 worker邻居加载时自动跨机拉取远端特征机制示意见 docs/source/_figures/dist_sampling.png训练入口参考 examples/multi_gpu/distributed_sampling.py多机提交脚本在同目录.sbatch文件里。点云任务上更常见的是单机多卡用标准DistributedDataParallelDistributedSampler即可痛点推理延迟高→ 解法模型走torch.jit.script导出静态图PyG 消息传递层均支持 TorchScript示例见 examples/jit/ 目录导出前先把SamplePoints之类的随机变换挪到离线预处理pre_transform推理期只做确定性计算痛点点云变长显存波动大→ 解法控制T.SamplePoints(n)的点数上限1024 是分类常用值分割需要更密k不必贪大KNN 建边是 O(nk) 级别的开销k16~30已够用批大小按单卡显存实测点云任务的 batch 往往比 CV 小示例里是 10~32痛点评估口径→ 解法分类看top-1 accmodel(data).max(dim1)[1]分割用jaccard_indexexamples/dgcnn_segmentation.py 的做法两者都按data.batch或data.num_nodes归一别按 batch 数归一避坑清单CUDA 扩展与 torch 版本不匹配torch-scatter/pyg-lib等 wheel 是按 torch 主版本编译的报错多为undefined symbol或 import 失败。解法对照 docs/source/get_started/ 里与你 torch 版本对应的安装源不要混装忘记传batch向量DynamicEdgeConv、fps、global_max_pool都要求batch参数。漏传后 KNN 会跨图建边——第 1 张点云的点被算成第 2 张的邻居loss 正常但 acc 异常低fps的ratio设置过激ratio0.25意味着点云逐层缩到 1/4若初始只采样了 256 个点两层之后剩 16 个点特征严重不足。分类任务建议首层 ≥1024 点ratio从 0.5 起步SamplePoints放在transform里导致每次访问结果不同采样是随机的训练时可以当增广但做确定性评估或缓存时要用pre_transform缓存后不再随机把点云pos当特征x用Data.pos是坐标、Data.x才承载特征DynamicEdgeConv的边权 MLP 输入维度按两两拼接算2 * in_channels只有一路输入是它自动复制的别在维度上少算一倍资源与延伸点云教程含 KNNGraph 建图与 PointNet 消息传递公式推导docs/source/tutorial/point_cloud.rst可直接运行的示例examples/dgcnn_classification.py、examples/dgcnn_segmentation.py、examples/point_transformer_classification.py点云基准在 benchmark/points/相关论文DGCNNDynamic Graph CNN for Learning on Point Clouds、Point Transformer、PointNet均为上述示例文件头部引用的原始工作下一步建议先按上面的最小示例在 ModelNet10 上把 DGCNN 复现出来记录收敛曲线再切换到 ModelNet40 复现同一模型随后把DynamicEdgeConv替换为PointTransformerConv参考 point_transformer 示例的TransformerBlock在相同 epoch、相同SamplePoints(1024)设置下对比 top-1 acc你会直观看到长程注意力带来的提升幅度及其显存代价。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考