jittor-pct-modelnet40

基于 Jittor 实现的 Point Cloud Transformer（PCT）点云分类项目，用于在 ModelNet40 数据集上完成三维点云物体分类训练与预测。本项目适配本地环境和 OpenI/C2Net 训练平台，包含数据读取、点云增强、局部邻域建模、Transformer 编码、训练调度、checkpoint 保存和测试集结果生成流程。

项目特点

• 基于 Jittor 编写，支持 CUDA 训练，CUDA 不可用时自动回退到 CPU 调试配置。
• 使用局部 KNN 邻域聚合和 PCT 自注意力编码提取点云特征。
• 默认使用 FPS anchor 采样提升局部覆盖质量，可切换为 fast 采样以提高速度。
• 支持 --config JSON 配置文件，命令行参数优先级高于配置文件。
• 默认划分 10% 训练数据作为验证集，用验证准确率选择 checkpoint，降低训练集准确率偏乐观带来的线上掉分。
• 自动保存 config.json、command.txt、train.log，便于复现实验。
• 支持 OpenI/C2Net 数据集挂载、输出目录适配和结果回传。

目录结构

jittor-pct-modelnet40/
├── configs/
│   └── default.json        # 默认训练配置
├── data/
│   └── README.md           # 数据准备说明，不提交大文件
├── scripts/
│   ├── train_local.sh      # 本地训练示例
│   └── train_openi.sh      # OpenI 训练示例
├── pct.py                  # 训练、评测和预测主程序
├── requirements.txt        # Python 依赖
├── README.md
├── LICENSE
└── .gitignore

环境安装

推荐环境：

Python 3.8+
Jittor 1.3+
NumPy
CUDA GPU 环境，推荐 A100 或同等级 GPU

安装依赖：

pip install -r requirements.txt

OpenI 平台通常已在 Jittor GPU 镜像中预装 Jittor 和 C2Net SDK。如果本地没有 CUDA，脚本会自动切换到较小 CPU 调试配置，但训练速度和最终结果会明显低于 GPU。

数据准备

将预处理后的 ModelNet40 数据放到 data/ 目录：

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

如果数据在其他路径，可通过 --data_dir 或 --data_root 指定：

python pct.py --data_dir /path/to/modelnet40_npy

.npy 数据文件默认不会提交到 Git。更多说明见 data/README.md。

训练

使用默认配置训练：

python pct.py --config configs/default.json

本地脚本：

bash scripts/train_local.sh

OpenI 脚本：

bash scripts/train_openi.sh

如果训练时间紧，可以使用快速 anchor 采样：

python pct.py --config configs/default.json --anchor_method fast

如果希望冲击更高线上分数，可以适当增加训练轮数和测试投票次数：

python pct.py --config configs/default.json --epochs 260 --vote_num 16

评测与推理

当前脚本在训练结束后会自动加载最优 checkpoint，并对测试集生成预测文件：

outputs/result.json

使用已有 checkpoint 直接推理：

python pct.py --eval_only --checkpoint outputs/pct_model_best.pkl --data_dir data

如果在 OpenI 平台运行，结果会保存到平台提供的 output 目录，并在训练任务结束时调用 upload_output() 回传。

CPU 小规模调试命令：

python pct.py --epochs 5 --batch_size 8 --n_points 512 --anchor_method fast

结果说明

训练日志中的 Train Acc 表示训练集分类准确率。若设置 --val_ratio 大于 0，脚本会从训练集中划分验证集并输出 Val Acc。线上最终成绩以平台根据 result.json 计算的测试集准确率为准。

默认配置偏向单卡 A100 上的精度和速度平衡，并使用验证集准确率选择最优 checkpoint。由于随机种子、平台环境、Jittor/CUDA 版本、测试投票次数、验证集划分和是否使用 FPS anchor 等因素不同，本地复现结果可能与线上成绩存在小幅差异。若需要使用全部训练数据训练，可以设置 --val_ratio 0 --selection_metric train，但该模式下 checkpoint 选择会更依赖训练准确率，线上分数可能更不稳定。

可复现记录

每次运行会在输出目录保存：

• config.json：本次实际使用的参数。
• command.txt：启动命令。
• train.log：训练日志。
• pct_model.pkl：最终模型。
• pct_model_best.pkl：最优模型。
• pct_model_top*.pkl：用于集成预测的 checkpoint。
• result.json：测试集预测结果。

这些文件默认被 .gitignore 忽略。

主要参数

参数	默认值	说明
--config	None	JSON 配置文件路径
--data_dir / --data_root	data	ModelNet40 数据目录
--output_dir	outputs	本地输出目录
--n_points	1024	每个样本采样点数
--batch_size	48	batch size
--epochs	240	训练轮数
--lr	8e-4	初始学习率
--vote_num	12	测试投票次数
--val_ratio	0.1	从训练集划分出的验证集比例
--selection_metric	auto	auto 优先用验证准确率选模型，无验证集时用训练准确率
--stage1_points	192	局部 anchor 点数量
--knn_k	24	KNN 邻域点数量
--anchor_method	fps	fps 精度更优，fast 速度更快
--seed	42	随机种子
--eval_only	False	只加载 checkpoint 推理，不训练
--checkpoint	None	推理模式下的模型路径

Git 忽略规则说明

.gitignore 会忽略 Python 缓存、Jittor 编译缓存、虚拟环境、数据集大文件、checkpoint、日志和输出结果。.gitignore 只影响未跟踪文件；如果大文件已被 Git 跟踪，需要先执行：

git rm --cached path/to/file

再提交更新后的忽略规则。

第三方引用

PCT 模型思想来自 Point Cloud Transformer 相关工作；本仓库代码基于 Jittor API 自行实现和适配。数据集为 ModelNet40，请遵循数据集原始许可和比赛平台要求。

贡献

欢迎提交 issue 或 pull request 改进模型结构、训练策略、OpenI 适配和文档说明。提交前请确认：

• 代码至少通过基本语法检查。
• README 中命令与当前代码一致。
• 不提交数据集大文件、缓存、checkpoint 和输出结果。

许可证

本项目采用 MIT License 开源，详见 LICENSE 文件。