IterativePFN-Jittor 点云降噪

本项目是第七届计图人工智能挑战赛赛道二的参赛实现。项目使用 Jittor 复现 IterativePFN，并针对比赛使用的 Laplace 噪声进行分阶段微调。最终模型采用动态 EdgeConv、四个内部迭代模块、自适应最近邻监督和反向覆盖损失。

环境安装

推荐环境：

Ubuntu 22.04/24.04
Python 3.10-3.12
Jittor 1.3.10
NVIDIA CUDA 12.x（也可使用 Jittor 自动安装的 CUDA）

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
python -m jittor.test.test_example
PYTHONPATH=. python tools/smoke_test.py --cuda

若系统存在多个编译器或 CUDA，可显式设置：

export cc_path=/usr/bin/g++-12
export nvcc_path=/usr/local/cuda/bin/nvcc

数据准备

比赛数据不包含在仓库中。训练集解压后应满足：

dataset_train/
└── shapenet/<category>/<model_id>/models/model_normalized.obj

dataset_test_noisy/
└── shapenet/<category>/<model_id>/noisy.npy

训练集可使用比赛提供的下载地址：

wget -O dataset_train.tar.gz \
  "https://cloud.tsinghua.edu.cn/f/73bc6ca7ac6344949b07/?dl=1"
tar -xzf dataset_train.tar.gz

测试集请从比赛页面下载。仓库内的 data/datalist/ 仅保存划分列表，不包含原始数据。

预训练权重

官方数据与 PyTorch 预训练模型由 IterativePFN 作者提供：

https://drive.google.com/file/d/1-TvHy3bvq8X1vI0ztwmmubDqhngRLQDu/view

下载并解压 data_and_ckpt.zip，然后转换权重：

pip install torch
python tools/convert_checkpoint.py \
  --input pretrained/denoisenet-ep-99.ckpt \
  --output pretrained/denoisenet-ep-99.npz

权重文件较大，按开源规范不提交到 Git。

训练

以下命令复现最终四阶段微调流程。每阶段会在输出目录保存 config.json、command.txt、train.log、best.npz 和 last.npz。

bash scripts/train_final.sh \
  ./dataset_train \
  ./pretrained/denoisenet-ep-99.npz \
  ./outputs/final

所有阶段均固定随机种子。显存不足时可在脚本中将 --patch-size 调小，但结果会与报告成绩存在差异。

推理

单卡推理：

PYTHONPATH=. python -m src.iterativepfn.infer \
  --input ./dataset_test_noisy \
  --output ./outputs/prediction \
  --checkpoint ./outputs/final/stage4/best.npz \
  --batch-size 2

四卡并行推理并打包：

bash scripts/infer_four_gpu.sh \
  ./dataset_test_noisy \
  ./outputs/final/stage4/best.npz \
  ./outputs/prediction

python tools/package_submission.py \
  --input ./outputs/prediction \
  --output ./outputs/result.zip

结果说明

比赛指标由 Chamfer Distance（CD）和 Point-to-Surface（P2S）组成：

score = 0.5 * CD_score + 0.5 * P2S_score
metric_score = clip(100 * (1 - metric_pred / metric_noisy), 0, 100)

本仓库对应方案的线上结果为：

Score	CD score	P2S score
75.53	约 64	约 87

线上评测使用隐藏 clean 点云和网格。本地验证通过公开训练网格重新采样得到，因此本地代理指标与线上分数不完全一致。GPU、Jittor/CUDA 版本和随机采样也可能造成小幅差异。

完整的算法、创新点和实验分析见 REPORT.md。

项目结构

configs/                 最终训练参数记录
data/datalist/           比赛数据划分列表
scripts/                 训练与四卡推理脚本
src/iterativepfn/        模型、训练和推理核心代码
tools/                   权重转换、打包和检查工具
outputs/                 运行产物（Git 忽略）

致谢

模型结构和官方预训练权重来自 IterativePFN，详情及许可证见 NOTICE。本项目的 Jittor 移植、Laplace 噪声微调、覆盖损失和比赛推理流程由参赛者实现。