ModelNet40 Point Cloud Classification (Jittor PCT)

本仓库是计图比赛 / 计算机图形学 PA3 的 ModelNet40 点云分类代码。核心模型为基于 Jittor 实现的 PCT（Point Cloud Transformer），最终结果可通过 seed=7777 的训练配置复现。

1. 环境安装

推荐环境：

Python 3.9
Jittor
NumPy
CUDA 环境可选；没有 GPU 时可将 --use_cuda 设为 0

安装依赖：

pip install -r requirements.txt

如果在 WSL 中运行 Jittor 且遇到 g++ / nvcc 兼容问题，可在命令前增加类似 env cc_path=/usr/bin/g++-12 的编译器设置，具体路径以本机环境为准。

2. 数据准备

本仓库不提交数据集原始文件。请将比赛提供的 ModelNet40 预处理数据放入 data/ 目录：

data/
├── categories.txt
├── train_points.npy
├── train_labels.npy
└── test_points.npy

其中 train_points.npy 和 test_points.npy 为点云数组，train_labels.npy 为训练标签。更多说明见 data/README.md。

3. 训练

最终单模型结果使用 seed=7777。推荐运行命令如下：

python pct.py \
  --data_dir ./data \
  --epochs 300 \
  --batch_size 32 \
  --n_points 2048 \
  --num_workers 4 \
  --optimizer adamw \
  --lr 0.0007 \
  --weight_decay 1e-4 \
  --seed 7777 \
  --use_cuda 1 \
  --log_interval 100 \
  --val_ratio 0 \
  --vote_num 10 \
  --tta_rotate 1 \
  --tta_scale_low 0.98 \
  --tta_scale_high 1.02 \
  --tta_jitter_std 0.005 \
  --tta_jitter_clip 0.01 \
  --save_model pct_s7777.pkl \
  --result_path result_s7777.json \
  --save_logits_path logits_s7777.npy

同样的参数已经记录在 configs/train_s7777.json。如果显存不足，可把 --batch_size 改为 16。

也可以使用脚本：

bash scripts/train.sh

4. 推理 / 生成提交文件

如果已经训练得到 pct_s7777.pkl，可直接加载模型生成测试集预测：

python pct.py \
  --data_dir ./data \
  --eval_only 1 \
  --load_model pct_s7777.pkl \
  --batch_size 32 \
  --n_points 2048 \
  --num_workers 4 \
  --seed 7777 \
  --use_cuda 1 \
  --vote_num 10 \
  --tta_rotate 1 \
  --result_path result.json

也可以使用脚本：

bash scripts/infer.sh
zip result.zip result.json

比赛平台提交 result.zip，其中只包含 result.json。GitLink / GitHub 开源仓库提交代码和复现说明，不提交 .npy 数据、.pkl 权重、日志或中间结果。

5. 结果说明

模型对测试集输出 2468 个样本的类别预测，结果格式为：

{
  "0": 12,
  "1": 3
}

由于训练随机性、Jittor / CUDA 版本和硬件差异，重新训练得到的线上分数可能与最终提交结果有小幅差异。复现时请优先保持 seed=7777、点数、训练轮数、学习率、TTA 参数与本文档一致。

6. 文件说明

pct.py：PCT 模型、训练、验证和推理主程序
ensemble_soft.py：可选的 logits 软融合脚本
search_ensemble_candidates.py：可选的融合候选搜索脚本
configs/train_s7777.json：最终 s_7777 实验配置记录
scripts/train.sh：训练入口脚本
scripts/infer.sh：推理入口脚本
data/README.md：数据准备说明

7. 许可证

本项目使用 MIT License，见 LICENSE。