PCT on ModelNet40 (Jittor 实现)

基于 Jittor 复现 Point Cloud Transformer (PCT) 在 ModelNet40 数据集上的点云分类任务。

论文：Guo et al., PCT: Point Cloud Transformer, Computational Visual Media 2021.

1. 项目结构

.
├── pct.py             主程序（含数据集、模型、训练、预测）
├── pct_model.pkl      训练好的模型权重（约 5.3 MB）
├── data/
│   └── categories.txt 40 个 ModelNet40 类别名
├── env.sh             环境变量旁路脚本（详见第 3 节）
├── requirements.txt   Python 依赖版本
├── LICENSE            MIT
└── README.md

数据集（train_points.npy / train_labels.npy / test_points.npy）未包含在本仓库中，请按第 4 节获取。

2. 模型与训练设置

项	值
网络	Input embedding (2× Conv1d) → 4× Self-Attention (Offset Attention, L1-norm) → Conv fuse 512→1024 → Global Max Pool → MLP 1024→512→256→40
优化器	SGD (momentum=0.9, weight_decay=1e-4)
学习率	初始 0.01，CosineAnnealing 衰减至 1e-3，T_max=250
Batch size	32
采样点数	1024
数据增强	绕 Y 轴随机旋转
训练轮数	250 epoch

最终训练精度：train_acc = **91.9%**（在单卡 RTX 4090 上跑 250 epoch，总耗时约 50 分钟，每 epoch ≈ 11.7s）。

3. 环境要求

• Linux + NVIDIA GPU（建议显存 ≥ 4 GB）
• NVIDIA driver ≥ 525（CUDA 12.x 兼容）
• Python 3.10
• jittor == 1.3.11.0
• numpy >= 1.26

3.1 安装 Jittor

推荐用 conda：

conda create -n jittor python=3.10 -y
conda activate jittor
pip install -r requirements.txt

Jittor 在首次 import 时会调用 nvcc 编译 CUDA 算子，需要一个与 GCC 兼容的 nvcc：

系统 nvcc	GCC	是否可用
≥ 12.0	11.x	可用
11.5（Ubuntu 22.04 自带）	11.x	不兼容，会报 parameter packs not expanded
≥ 11.6	10.x / 9.x	可用

如系统 nvcc 与 GCC 不兼容，最简单的办法是通过 conda 装一份新版 nvcc：

conda install -c conda-forge cuda-nvcc cuda-cudart cuda-cudart-dev cuda-libraries-dev cuda-version=12.1
conda install -c conda-forge cudnn=8.9.7   # Jittor 1.3.11 需要 cuDNN 8.x

注意：Jittor 1.3.11 调用的是 cuDNN 8 API，cuDNN 9 的接口变动会导致编译失败，请勿安装 cuDNN 9.x。

3.2 国内/校园网下载

Jittor 在初始化时会从 GitHub 下载 NCCL、cutt、cutlass、MKL、MPI 等第三方源码并就地编译；如果你的网络无法访问 GitHub，可通过本仓库提供的 env.sh 一键关闭这些非必要依赖：

source env.sh         # 等同于 export use_nccl=0 use_cutt=0 use_cutlass=0 use_mkl=0 use_mpi=0
python pct.py ...

PCT 在单卡训练下不需要这些库，关闭后不影响精度。

如果 conda-forge 的 CUDA 头文件只装到了 <env>/targets/x86_64-linux/include/，Jittor 可能找不到 cusparse.h；执行下面的 symlink 即可：

ENV=$(conda info --base)/envs/jittor
cd "$ENV/include"
for h in $(ls ../targets/x86_64-linux/include/); do
  [ ! -e "$h" ] && ln -sf "../targets/x86_64-linux/include/$h" "$h"
done
[ ! -e "$ENV/lib64" ] && ln -s lib "$ENV/lib64"

4. 数据获取

本仓库不分发数据，请按以下任一方式获取：

• 课程参与者：从课程指定渠道下载预切好的 train_points.npy / train_labels.npy / test_points.npy，放入 data/ 目录。
• 公开渠道：原始 ModelNet40 数据集来自 Princeton（官方页面）；如需自行预处理为同形态的 npy 文件，请参考 PCT 论文与官方代码。

文件预期形状（参见 pct.py 中 ModelNet40Dataset）：

文件	形状	说明
train_points.npy	(N_train, P, 3) float32	训练点云
train_labels.npy	(N_train,) int64	类别索引（0–39）
test_points.npy	(N_test, P, 3) float32	测试点云
data/categories.txt	40 行	类别索引到名称的映射

P 为每个样本预存的点数（默认从中随机采 --num_points 个）。

5. 复现

5.1 从零训练

source env.sh                             # 仅国内网络需要
python pct.py --use_cuda 1 \
              --epochs 250 \
              --batch_size 32 \
              --num_points 1024 \
              --save_model pct_model.pkl \
              --save_result result.json

完成后会生成：

• pct_model.pkl：训练好的模型权重
• result.json：测试集预测，格式 {"0": 17, "1": 3, ...}（key 为测试样本下标，value 为类别索引）

5.2 仅用提供的权重预测

直接复用本仓库的 pct_model.pkl 跑预测（需要先准备好 data/test_points.npy，并自行写一个加载脚本或对 pct.py 做轻微改动以跳过训练阶段）。

5.3 快速烟雾测试

source env.sh
python pct.py --use_cuda 1 --dry_run    # 2 epoch + 小 batch + 256 点

6. 致谢

• 论文：Guo M.-H., Cai J.-X., Liu Z.-N., Mu T.-J., Martin R. R., Hu S.-M. PCT: Point Cloud Transformer. Computational Visual Media, 2021.
• 框架：Jittor (https://github.com/Jittor/jittor)
• 数据：Princeton ModelNet40

7. License

MIT