点云降噪赛题 Baseline

环境安装

# 安装计图
conda create -n jittor python=3.9 -y
conda activate jittor
conda install -c conda-forge gcc=10 gxx=10 -y # 确保gcc、g++版本不高于10
conda install -c conda-forge libgomp -y # 确保OpenMP runtime存在

# 安装依赖
python -m pip install -r requirements.txt
pip install jittor numpy trimesh scipy omegaconf point-cloud-utils

数据准备

将训练数据 dataset_train.tar.gz 解压到本目录下：
```
tar xzf dataset_train.tar.gz
```
解压后目录：dataset_train/shapenet/<synset_id>/<model_id>/models/model_normalized.obj
将测试数据 dataset_test_noisy.zip 解压到本目录下：
```
unzip dataset_test_noisy.zip
```
解压后目录：dataset_test_noisy/shapenet/<synset_id>/<model_id>/noisy.npy

训练

bash run_with_log.sh jittor_train \
  python run.py --task configs/task/train_vm.yaml

训练权重保存在 experiments/ 目录下。

PyTorch Diffusion 版本

项目中另外新增了一套独立的 PyTorch 点云去噪实现，代码位于 torch_diffusion/，不会覆盖现有 Jittor baseline。

模型说明

条件扩散：学习 clean patch | noisy patch
骨干网络：Point Transformer 风格的局部 KNN 注意力网络
训练目标：DDPM epsilon prediction
推理方式：DDIM 采样 + patch 聚合回整点云

训练（2 张 4090）

bash run_with_log.sh torch_train_2gpu \
  bash launch_torch.sh --nproc_per_node=2 train_torch_diffusion.py \
  --config configs_torch/train_point_diffusion.yaml

如果只是做单卡 smoke test，不需要 torchrun，可以直接运行：

bash run_with_log.sh torch_train_1gpu \
  python train_torch_diffusion.py \
  --config configs_torch/train_point_diffusion_fast.yaml

默认训练配置：

patch_size=1024
batch_size=24（每卡）
bf16 混合精度
hidden_dim=256
depth=8

checkpoint 会保存在 experiments_torch/point_diffusion/。

如果你更在意训练速度，可以直接用更激进的快速配置：

bash run_with_log.sh torch_train_fast \
  bash launch_torch.sh --nproc_per_node=2 train_torch_diffusion.py \
  --config configs_torch/train_point_diffusion_fast.yaml

这版会进一步缩小模型、patch 和采样步数，用更短时间先训出一版可用模型。

如果服务器是共享环境，或者你更在意“更快出一个能看的 checkpoint”而不是单次配置的上限吞吐，可以直接使用：

bash run_with_log.sh torch_train_server_fast \
  python train_torch_diffusion.py \
  --config configs_torch/train_point_diffusion_server_fast.yaml

这版额外做了几件事：

用更轻的骨干配置（hidden_dim=144, depth=3, knn_k=12）
将扩散步数降到 200
将验证集样本数降到 16，并把验证采样步数降到 12
用 total_num_workers 控制总 CPU worker 数，减少共享服务器上的抢占

如果你明确想用双卡提速，并且希望每 5 个 epoch 就做一轮验证，可以使用：

bash run_with_log.sh torch_train_2gpu_eval5 \
  bash launch_torch.sh --auto-gpus 2 --nproc_per_node=2 train_torch_diffusion.py \
  --config configs_torch/train_point_diffusion_2gpu_eval5.yaml

这版的取舍是：

保留双卡 DDP 加速
eval_interval=5
进一步缩小了 backbone、patch 和 diffusion 步数
通过 max_shapes=2048 控制每个 epoch 的训练样本量，优先把单 epoch 压短
更适合作为“先看趋势、再决定要不要冲更高分”的训练配置

launch_torch.sh 现在支持 --auto-gpus N，会在没有手动设置 CUDA_VISIBLE_DEVICES 时，按“空闲显存更多、已用显存更少、利用率更低”的顺序自动挑卡。

如果服务器卡比较多，也可以直接使用 4 卡高吞吐配置：

bash run_with_log.sh torch_train_4gpu \
  env OMP_NUM_THREADS=4 bash launch_torch.sh --nproc_per_node=4 train_torch_diffusion.py \
  --config configs_torch/train_point_diffusion_4gpu.yaml

这版额外启用了更高的 dataloader 并发、mesh 缓存，以及可选 torch.compile。

当前 4gpu 配置已经按“总共约 4 个 dataloader CPU workers”做了收敛，避免 4 卡时把 num_workers 按卡数放大到几十个后台进程。

如果日志里出现 env: 'torchrun': No such file or directory，说明当前环境里找不到 torchrun 可执行文件；如果同时 python -c "import torch" 也失败，那就是 PyTorch 本体还没装好，需要先切到正确环境或重新安装 PyTorch。

推理

先把 configs_torch/predict_point_diffusion.yaml 里的 load_ckpt 改成你的权重路径，然后运行：

bash run_with_log.sh torch_predict \
  python predict_torch_diffusion.py \
  --config configs_torch/predict_point_diffusion.yaml

输出保存在 results_torch/dataset_test_noisy/，目录结构与测试集保持一致，文件名为 denoised.npy。

推理（生成提交文件）

修改 configs/task/predict_vm.yaml 中的 load_ckpt 为你的最佳权重路径，然后运行：

bash run_with_log.sh jittor_predict \
  python run.py --task configs/task/predict_vm.yaml

降噪结果保存在 results/ 目录下，格式为 .npy (float32, shape (N,3))。

后台运行与日志

如果服务器不稳定，建议统一通过 run_with_log.sh 启动任务。这个脚本会：

自动将标准输出和错误输出写入 logs/<任务名>_时间戳.log
自动记录 PID 到 logs/<任务名>.pid
优先使用 nohup + setsid，在终端断开后继续运行

常用查看方式：

cat logs/torch_train_2gpu.pid

上面命令用于查看后台任务 PID。确认进程是否还在运行可以执行：

ps -fp "$(cat logs/torch_train_2gpu.pid)"

查看训练日志请使用：

tail -f logs/torch_train_2gpu_*.log

如果要停止任务，可以执行：

kill "$(cat logs/torch_train_2gpu.pid)"

打包提交

cd results/dataset_test_noisy
zip -r ../../result.zip shapenet/

提交格式

每个测试样本一个 denoised.npy，目录结构与测试集一致，打包为 result.zip：

result.zip
  shapenet/
    <synset_id>/
      <model_id>/
        denoised.npy    # np.float32, shape (N, 3)

本地评测（需要 GT 数据，仅组委会持有）

python evaluate.py \
    --pred_dir ./results/dataset_test_noisy \
    --gt_dir ./test_gt \
    --noisy_dir ./dataset_test_noisy \
    --mesh_dir ./dataset_train \
    --workers 8