1. 环境配置:为RTX 4090搭建专属的PointGroup运行环境

想在RTX 4090上跑通PointGroup,第一步也是最关键的一步,就是搭建一个“严丝合缝”的软件环境。这就像给一台顶级跑车加对型号的汽油,油不对,再强的性能也发挥不出来。我刚开始折腾的时候,也踩了不少坑,特别是那个经典的 cublas runtime error,相信很多用40系显卡的朋友都遇到过。今天我就把最稳、最省事的配置流程分享给你。

1.1 核心依赖版本选择:为什么是Python 3.7 + CUDA 11.6 + Torch 1.13.1?

原项目的README可能推荐的是比较老的Torch版本,但直接用在RTX 4090上大概率会出问题。这是因为RTX 4090基于Ada Lovelace架构,需要CUDA 11.6或更高版本才能充分发挥其性能,并且与特定版本的PyTorch有兼容性要求。

我实测下来,Python 3.7CUDA 11.6 搭配 PyTorch 1.13.1 这个组合在RTX 4090上是最稳定的。Python 3.7是一个比较经典的版本,很多老项目的依赖都对其有良好支持。CUDA 11.6是官方支持RTX 40系列显卡的较早版本,既能保证兼容性,又拥有完善的生态。PyTorch 1.13.1是最后一个官方明确支持CUDA 11.6的稳定版本之一,后续的版本可能默认支持更新的CUDA,手动编译或找对应版本会更麻烦。

所以,第一步就是创建并激活我们的专属环境:

conda create -n pointgroup python=3.7 -y
conda activate pointgroup

1.2 一步步安装PyTorch与基础依赖

激活环境后,我们使用PyTorch官方提供的针对CUDA 11.6的预编译版本来安装。下面的命令会一次性安装PyTorch、Torchvision和Torchaudio,注意后面的 --extra-index-url 参数必须带上,它指定了下载这些特定版本wheel包的地址。

pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116

安装完成后,强烈建议你验证一下CUDA是否真的可用。打开Python解释器,运行:

import torch
print(torch.__version__)  # 应该输出 1.13.1+cu116
print(torch.cuda.is_available())  # 应该输出 True
print(torch.cuda.get_device_name(0))  # 应该显示你的RTX 4090

接下来,安装项目运行所需的其他基础Python包。这些是数据处理、编译和日志记录所必需的:

pip install cmake plyfile tensorboardX pyyaml scipy

1.3 编译安装spconv:解决CUDNN路径与算子冲突

PointGroup依赖一个名为 spconv 的库来进行高效的稀疏卷积运算,这个库通常需要从源码编译。这是整个环境搭建中最容易出错的一环,但别担心,跟着我的步骤来。

首先,我们需要安装两个系统级的依赖,它们能通过Conda方便地获取:

conda install libboost -y
conda install -c bioconda google-sparsehash -y

然后,进入项目中的 lib/spconv 目录进行编译:

cd lib/spconv
python setup.py bdist_wheel

坑点一:找不到CUDNN。 编译时可能会报错,提示找不到cuDNN库。这是因为CMake没有自动定位到你的cuDNN路径。解决方法是在 lib/spconv/CMakeLists.txt 文件的开头部分(比如在 project(SpConv) 语句之后),手动添加以下四行,将路径替换成你自己环境中cudnn的实际安装路径。你可以通过 conda list | grep cudnn 来查找。

set(CUDNN_INCLUDE_DIR "/home/你的用户名/.conda/pkgs/cudnn-8.2.1.32-h86fa8c9_0/include")
set(CUDNN_INCLUDE_PATH "/home/你的用户名/.conda/pkgs/cudnn-8.2.1.32-h86fa8c9_0/include")
set(CUDNN_LIBRARY "/home/你的用户名/.conda/pkgs/cudnn-8.2.1.32-h86fa8c9_0/lib/libcudnn.so")
set(CUDNN_LIBRARY_PATH "/home/你的用户名/.conda/pkgs/cudnn-8.2.1.32-h86fa8c9_0/lib/libcudnn.so")

坑点二:API变更错误。 由于我们使用了较新的Torch 1.13.1,可能会遇到 torch::jit::RegisterOperators 相关的编译错误。你需要找到报错对应的源文件(通常是以 .cpp 结尾的文件),将 torch::jit::RegisterOperators() 替换为 torch::RegisterOperators()。这是PyTorch版本升级导致的API变化。

坑点三:不支持的GPU算力。 RTX 4090的算力(Compute Capability)是8.9。有时编译脚本中预设的算力列表可能不包含8.9,或者CUDA工具链默认不支持。一个简单的解决方法是临时降低目标算力。在编译前,在终端中执行:

export TORCH_CUDA_ARCH_LIST="8.6"

这会让编译器以为是为算力8.6的显卡(如RTX 30系列)编译,但生成的代码在8.9的4090上通常是兼容且能运行的。这是一种实用的变通方法。

坑点四:运算符重载冲突。 可能会遇到 more than one operator ">" matches these operands 这类错误。这是因为CUDA的头文件和PyTorch的头文件定义了相同的运算符,编译器懵了。解决方法是在刚才的 CMakeLists.txt 里再添加一个编译定义:

add_definitions(-D__CUDA_NO_HALF_OPERATORS__)

解决完上述问题后,重新运行 python setup.py bdist_wheel。编译成功后,会在 dist 文件夹下生成一个 .whl 文件,用pip安装它即可:

cd dist
pip install spconv-1.0-cp37-cp37m-linux_x86_64.whl

1.4 编译PointGroup自定义算子

spconv搞定后,还需要编译PointGroup项目自身的自定义C++/CUDA算子。回到项目根目录,进入 lib/pointgroup_ops

cd ../../../lib/pointgroup_ops
python setup.py develop

这里可能遇到的典型错误是 fatal error: THC/THC.h: No such file or directory。这是因为在PyTorch 1.13中,THC/THC.h 这个头文件可能已经被移除或改变了位置。解决方法通常是打开报错的文件(如 src/bfs_cluster/bfs_cluster.h),找到 #include <THC/THC.h> 这一行,直接将其删除。这个头文件在较新版本的PyTorch中往往已经不再需要。

如果一切顺利,develop 模式安装成功,你的核心环境就配置完成了。恭喜你,最复杂的部分已经过去!

2. 数据准备:处理ScanNet v2数据集并理解其结构

模型有了,接下来就是“喂”给它的数据。PointGroup最常用的数据集是ScanNet v2,这是一个大规模的室内3D场景数据集。处理数据的过程有点繁琐,但理解了这个流程,你以后处理任何自定义点云数据都会得心应手。

2.1 下载与初步组织

首先,你需要从ScanNet官网申请下载数据。下载后,你会得到类似 scans_trainscans_test 等文件夹。原始结构比较分散,我们需要将其整理成PointGroup代码期望的格式。假设你的数据集根目录是 dataset/scannetv2,下面是一个我常用的Python整理脚本,它能帮你把文件复制到正确的位置:

import os
import shutil

root_dir = "dataset/scannetv2"
dataset_dir_src = "dataset/scannetv2/scans"  # 假设这是解压后的主文件夹
dataset_train = os.path.join(root_dir, 'train')
dataset_val = os.path.join(root_dir, 'val')
dataset_test = os.path.join(root_dir, 'test')

# 我们需要的文件类型后缀
sub_str_list = ['_vh_clean_2.ply', '_vh_clean_2.labels.ply', '_vh_clean_2.0.010000.segs.json', '.aggregation.json']

# 创建目标文件夹
for d in [dataset_train, dataset_val, dataset_test]:
    os.makedirs(d, exist_ok=True)

# 复制训练集文件 (ScanNet的train split)
for root, dirs, files in os.walk(os.path.join(dataset_dir_src, 'scans_train')):
    for file in files:
        for sub_str in sub_str_list:
            if sub_str in file:
                shutil.copy(os.path.join(root, file), os.path.join(dataset_train, file))
                break

# 复制验证集文件 (通常从训练集中划分一部分作为val,这里假设你已经知道哪些场景是val)
# 你需要一个包含验证集场景名的列表,例如 val_scenes.txt
val_scene_list = [...] # 这里填入你的验证集场景ID列表
for scene in val_scene_list:
    scene_path = os.path.join(dataset_dir_src, 'scans_train', scene)
    if os.path.exists(scene_path):
        for file in os.listdir(scene_path):
            for sub_str in sub_str_list:
                if sub_str in file:
                    shutil.copy(os.path.join(scene_path, file), os.path.join(dataset_val, file))

# 复制测试集文件
for root, dirs, files in os.walk(os.path.join(dataset_dir_src, 'scans_test')):
    for file in files:
        for sub_str in sub_str_list:
            if sub_str in file:
                shutil.copy(os.path.join(root, file), os.path.join(dataset_test, file))
                break

# 复制标签定义文件
shutil.copy(os.path.join(dataset_dir_src, 'scannetv2-labels.combined.tsv'),
            os.path.join(root_dir, 'scannetv2-labels.combined.tsv'))

2.2 运行官方预处理脚本生成训练数据

文件整理好后,我们需要运行PointGroup自带的预处理脚本,将原始的.ply.json文件转换成PyTorch可以直接加载的.pth文件。这个脚本会计算法线、颜色归一化、生成实例标签等。

cd dataset/scannetv2
python prepare_data_inst.py --data_split train
python prepare_data_inst.py --data_split val
python prepare_data_inst.py --data_split test

运行成功后,在 trainvaltest 文件夹下,你会看到多出了很多 *_inst_nostuff.pth 文件。这就是我们模型真正要读取的数据。我强烈建议你写个小脚本打开一个看看它的结构,理解数据格式对后续调试至关重要:

import torch
data = torch.load('dataset/scannetv2/train/scene0000_00_inst_nostuff.pth')
print(type(data)) # 应该是tuple
print(len(data))  # 应该是4
coord, color, semantic_label, instance_label = data
print(f"坐标形状: {coord.shape}")      # (N, 3)
print(f"颜色形状: {color.shape}")      # (N, 3),值域[-1,1]
print(f"语义标签形状: {semantic_label.shape}") # (N, 1),-100表示忽略
print(f"实例标签形状: {instance_label.shape}") # (N, 1),-100表示忽略

关键理解-100 在这里是一个“忽略”标记。在计算损失时,标签为-100的点会被自动忽略。这对于处理背景点(如墙壁、地板)或未标注的点非常有用。在PointGroup中,通常将语义标签为0(墙)和1(地板)的点作为背景,不参与实例聚类。

3. 模型训练与调参:在RTX 4090上启动你的第一次训练

环境齐备,数据就绪,终于到了激动人心的训练环节。RTX 4090的24GB大显存在这里优势尽显,你可以使用更大的批次大小(batch size)或更高分辨率的输入,从而可能获得更快的收敛速度和更好的模型性能。

3.1 理解配置文件与启动训练

PointGroup使用YAML配置文件来管理所有超参数。核心配置文件通常位于 config/ 目录下,例如 pointgroup_run1_scannet.yaml。在启动训练前,花几分钟浏览一下这个文件,里面定义了网络结构、优化器、学习率、数据路径等关键信息。

使用以下命令开始训练:

CUDA_VISIBLE_DEVICES=0 python train.py --config config/pointgroup_run1_scannet.yaml

CUDA_VISIBLE_DEVICES=0 指定使用第一块GPU(你的RTX 4090)。如果你有多块卡,可以尝试数据并行训练,但对于单机来说,一块4090通常已经足够强大。

训练开始后,你应该在终端看到损失值在不断下降。同时,项目会使用Tensorboard来记录训练过程。你可以用以下命令启动Tensorboard来实时监控训练曲线:

tensorboard --logdir exp/scannetv2/pointgroup/pointgroup_run1_scannet/tensorboard

然后在浏览器中打开 localhost:6006,查看损失曲线、学习率变化等,这对于判断模型是否在正常学习、是否过拟合至关重要。

3.2 针对RTX 4090与自定义数据的调参技巧

虽然默认配置在ScanNet上效果不错,但如果你想在自己的数据上训练,或者想榨干4090的性能,调参是必不可少的。这里分享几个我实践过的关键点:

1. 体素大小 (voxel_size): 这是最重要的参数之一,直接决定了计算量和精度。值越小,点云被离散化的程度越低,保留的细节越多,但显存消耗和计算量呈立方级增长。对于室内场景(如ScanNet),0.02米或0.03米是常用值。对于室外大场景(如自动驾驶点云),可能需要0.05米甚至更大。RTX 4090的大显存允许你尝试比论文中更小的体素(如0.015米),可能会提升对小物体的分割精度。

2. 聚类半径 (cluster_radius): 在PointGroup的后处理聚类阶段,这个半径决定了多大距离内的点会被归为同一个实例。对于家具密集的场景,这个值要设小一些(如0.03米),避免把两个紧挨着的椅子合并。对于空旷场景下的大物体,可以适当调大。这需要根据你数据集中物体的平均尺寸来调整。

3. 批次大小 (batch_size):config 文件的 DATA 部分可以找到。4090的24GB显存让你有很大的发挥空间。对于ScanNet,你可以尝试将batch_size从默认的4或6提高到8甚至12。更大的batch size通常能使梯度估计更稳定,有助于模型收敛,有时还能允许你使用稍大一点的基础学习率。

4. 学习率与优化器: 默认使用Adam优化器,学习率0.001是个不错的起点。如果你增大了batch_size,一个经验法则是可以按 new_lr = old_lr * sqrt(new_bs/old_bs) 的比例适当增大学习率。同时,关注Tensorboard中的训练损失曲线,如果损失在初期剧烈震荡然后下降很慢,可能是学习率太大了;如果下降极其缓慢,则可能是学习率太小。

5. 处理训练损失为负的情况: 如果你在训练中发现损失值变成了负数,先别慌。PointGroup的总损失由语义分割损失、偏移量回归损失等几部分组成。其中偏移量回归损失(L1 loss)在某些实现下,如果预测得非常准确,可能会使整体损失出现负值。只要验证集指标在正常提升,这就不是问题。

4. 模型推理与结果可视化:看看你的模型“看到了”什么

模型训练好后,我们自然要看看它的分割效果。PointGroup的推理流程包括在验证集上评估指标和在测试集上生成可视化结果。

4.1 验证集评估与指标解读

首先,我们需要为验证集生成ground truth的文本格式文件,用于计算评估指标:

cd dataset/scannetv2
python prepare_data_inst_gttxt.py
cd ../..

然后,使用训练好的权重进行推理。假设你的最佳模型权重保存在 checkpoint/pointgroup.pth

CUDA_VISIBLE_DEVICES=0 python test.py --config config/pointgroup_default_scannet.yaml --pretrain checkpoint/pointgroup.pth

运行结束后,你会看到一整套评估指标,其中最重要的是 mAP@0.25mAP@0.5。它们表示在IoU(交并比)阈值分别为0.25和0.5时的平均精度。通常,mAP@0.5是更严格的指标。ScanNet v2上优秀的模型mAP@0.5能达到35%以上,mAP@0.25能达到55%以上。这些指标能客观地告诉你模型的性能水平。

4.2 测试集推理与3D可视化

在验证集上确认模型效果不错后,就可以在测试集上运行并生成可视化的分割结果了。首先,你需要修改配置文件中的TEST部分,或者直接在命令行覆盖相关参数。更简单的方法是运行:

CUDA_VISIBLE_DEVICES=0 python test.py --config config/pointgroup_run1_scannet.yaml --pretrain checkpoint/pointgroup.pth --test_split test --save_instance

--save_instance 参数会告诉程序将预测的每个实例保存为单独的 .txt 文件,里面包含属于该实例的点的索引。

得到这些 .txt 文件后,原始的预测结果并不方便直观查看。我们需要将其转换为一个完整的、带颜色的点云文件。下面这个脚本是我常用的,它会把每个实例分配一个随机颜色,并生成一个可以用 CloudCompareMeshLab 打开的文本格式点云文件:

import numpy as np
import torch
import os

# 路径设置
points_path_root = "dataset/scannetv2/test"  # 测试集原始数据pth文件位置
pred_files_root = "exp/scannetv2/pointgroup/pointgroup_run1_scannet/result/epoch384_nmst0.3_scoret0.09_npointt100/test"  # 预测结果文件夹
save_path_root = os.path.join(pred_files_root, 'colored_clouds')
os.makedirs(save_path_root, exist_ok=True)

# 获取所有预测的场景文件
pred_scene_files = [f for f in os.listdir(pred_files_root) if f.endswith('.txt') and 'pred' in f] # 假设预测文件以‘pred’开头

for pred_file in pred_scene_files:
    scene_id = pred_file.replace('_pred.txt', '')
    # 加载对应的点云坐标和颜色
    pth_path = os.path.join(points_path_root, f'{scene_id}_inst_nostuff.pth')
    coord, color, _, _ = torch.load(pth_path) # 只取坐标和颜色
    coord = coord.numpy()
    color = (color.numpy() * 127.5 + 127.5).astype(np.uint8) # 将归一化颜色[-1,1]转回[0,255]

    # 初始化实例标签数组
    instance_labels = np.zeros(coord.shape[0], dtype=np.int32)

    # 读取预测文件,里面每一行是一个实例文件路径
    with open(os.path.join(pred_files_root, pred_file), 'r') as f:
        instance_paths = f.read().strip().splitlines()

    # 为每个实例分配一个ID并标记点
    for inst_id, inst_rel_path in enumerate(instance_paths, start=1):
        inst_abs_path = os.path.join(pred_files_root, inst_rel_path)
        if os.path.exists(inst_abs_path):
            # 假设实例文件里存储的是属于该实例的点的全局索引
            point_indices = np.loadtxt(inst_abs_path, dtype=np.int64)
            instance_labels[point_indices] = inst_id

    # 生成随机颜色映射(为每个实例ID生成一个RGB颜色)
    unique_labels = np.unique(instance_labels)
    color_map = {}
    for label in unique_labels:
        if label == 0: # 0表示背景或无实例
            color_map[label] = np.array([128, 128, 128]) # 灰色
        else:
            color_map[label] = np.random.randint(0, 256, size=3)

    # 根据实例标签为点云上色
    instance_colors = np.array([color_map[l] for l in instance_labels])

    # 保存为PLY格式(更通用)或XYZRGB格式
    save_path = os.path.join(save_path_root, f'{scene_id}_colored.ply')
    # 这里以简单的PLY格式为例
    with open(save_path, 'w') as f:
        f.write('ply\n')
        f.write('format ascii 1.0\n')
        f.write(f'element vertex {coord.shape[0]}\n')
        f.write('property float x\n')
        f.write('property float y\n')
        f.write('property float z\n')
        f.write('property uchar red\n')
        f.write('property uchar green\n')
        f.write('property uchar blue\n')
        f.write('end_header\n')
        for i in range(coord.shape[0]):
            f.write(f'{coord[i,0]} {coord[i,1]} {coord[i,2]} {instance_colors[i,0]} {instance_colors[i,1]} {instance_colors[i,2]}\n')
    print(f'场景 {scene_id} 可视化文件已保存至: {save_path}')

将生成的点云文件导入CloudCompare,你就能清晰地看到模型是如何将场景中的桌子、椅子、沙发等物体一个个区分开来的。这种直观的反馈,是调优模型和解决问题的强大工具。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐