PCT_jittor

A Jittor implementation of Point Cloud Transformer (PCT) for ModelNet40 point cloud classification.

基于 Jittor 实现的 Point Cloud Transformer（PCT），完成 ModelNet40 点云 40 分类任务（计图挑战赛赛道二热身赛）。

项目简介

动机与目标

点云是三维场景的重要表示形式。本项目的目的是在 ModelNet40 数据集上实现并训练 Point Cloud Transformer（PCT），对点云进行物体分类，并生成符合评测要求的预测文件 result.zip。

解决的问题

将预处理好的 ModelNet40 点云（.npy）输入 PCT 网络完成分类；
在 Jittor 框架下完成训练、验证与测试推理；
输出 2468 条测试样本的类别预测（result.json）。

项目特点

完整实现 PCT 主干：4 层 Self-Attention、多尺度特征融合、全局池化与全连接分类头；
训练阶段：Y 轴取反、点云归一化、随机几何增强、SGD + 余弦退火学习率；
推理阶段：测试时增强（TTA，多次随机采点投票）提升测试准确率；
按训练集 Clean-Acc（无增强）保存最佳模型权重。

学到什么

使用 Jittor 搭建点云深度学习流水线（Dataset、训练循环、模型保存）；
理解 PCT 中 offset-attention 与点云数据增强对泛化的影响；
区分「增强训练准确率」与「评测准确率」的差异，并通过 TTA 与归一化改进结果。

环境要求与安装

环境

Python 3.8+
NVIDIA GPU + CUDA（推荐），或 CPU（WSL 下可用 --cpu 兜底）

安装依赖

git clone https://www.gitlink.org.cn/NieZinan/jittor.git
cd jittor
pip install -r requirements.txt

数据准备

数据集体积较大，不包含在仓库中（见 .gitignore）。请从课程/比赛提供的 data.zip 解压到项目根目录：

unzip data.zip -d .

解压后应包含：

文件	形状	说明
`data/train_points.npy`	(9843, 2048, 3)	训练点云
`data/train_labels.npy`	(9843,)	训练标签（0–39）
`data/test_points.npy`	(2468, 2048, 3)	测试点云（无标签）

使用方法

1. GPU 训练（推荐）

./run_train_gpu.sh ./data 250

训练结束后自动生成 pct_model.pkl，并打包 result.zip（默认 TTA 15 票）。

2. 仅推理（已有模型）

./run_predict.sh 20

使用已有权重生成 result.zip（TTA 20 票，约 5–10 分钟）。

3. 命令行参数示例

# 完整训练
python3 pct.py --data_dir ./data --epochs 250 --batch_size 32 --lr 0.01 --pack --tta_votes 15

# 仅预测
python3 pct.py --data_dir ./data --predict_only --pack --tta_votes 20

# 从 checkpoint 微调
python3 pct.py --data_dir ./data --resume --epochs 80 --pack --tta_votes 15

4. WSL 下 GPU 报错时（CPU 训练）

若出现 CURAND_STATUS_INITIALIZATION_FAILED：

python3 pct.py --data_dir ./data --epochs 200 --pack --cpu

仓库结构

jittor/
├── pct.py              # 数据集、PCT 模型、训练与推理
├── run_train_gpu.sh    # GPU 训练脚本
├── run_train.sh        # CPU 训练脚本
├── run_predict.sh      # TTA 推理并打包
├── requirements.txt    # Python 依赖
├── README.md
└── .gitignore

实现说明

模块	说明
`ModelNet40Dataset`	加载 npy；训练时增强；测试时 Y 翻转 + 归一化 + 采点
`PCT`	Conv1d 嵌入 → 4×SA_Layer → concat → fuse → max pool → FC
训练	`zero_grad` + 梯度裁剪；按 Clean-Acc 保存最佳权重
推理	`predict_tta`：多次随机采点，对 logits 投票

评测提交

向 Jittor 评测平台提交 **result.zip**（根目录仅含 result.json）；
JSON 格式：{"0": 类别, "1": 类别, ...}，共 2468 条，类别编号 0–39；
通过线：测试集 Accuracy ≥ 0.80。

作者

姓名：聂子楠
学号：2024010672
开源仓库：https://www.gitlink.org.cn/NieZinan/jittor

许可证

本项目为课程作业开源代码，仅供学习与交流使用。

参考

README 写作指南：如何为你的 GitHub 项目写 README 文件（freeCodeCamp 中文）
Jittor 框架：https://github.com/Jittor/jittor