PCT-jittor

基于 Jittor 实现的 Point Cloud Transformer (PCT) 点云分类项目，用于 PA3 / ModelNet40 分类任务。代码完成数据加载、训练、模型保存、测试集预测和 result.json 导出。

目录结构

.
├── configs/              # 默认参数示例
├── data/                 # 数据目录说明，真实数据不提交
├── scripts/              # 训练和预测脚本
├── pct.py                # 核心模型、训练、预测代码
├── requirements.txt      # Python 依赖
├── LICENSE
├── .gitignore
└── README.md

生成文件不会提交到仓库，包括 data/*.npy、*.pkl、result.json、result.zip、 outputs/、logs/ 和 Jittor 缓存。

环境安装

推荐 Python 3.10。为了避免污染服务器环境，建议使用项目本地虚拟环境：

python3 -m venv .venv
.venv/bin/pip install -r requirements.txt

如果在 CUDA 服务器上运行，需确保系统已有 NVIDIA driver、CUDA toolkit 和可用的 nvcc。本项目会把 Jittor 缓存放在仓库内 .jittor_home/，避免写入全局环境。

数据准备

将 PA3 提供的数据放入 data/：

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

默认数据目录是 data，也可以通过 --data-dir 指定其他路径。程序在缺少数据文件时会给出明确报错和修复提示。

训练

CPU 训练可用于小样本 smoke test：

.venv/bin/python pct.py --mode train --cpu --data-dir data --model-path pct_model.pkl --epochs 1 --limit-train 64

CUDA 训练命令：

CUDA_VISIBLE_DEVICES=0 cuda_archs=80 .venv/bin/python pct.py --mode train --cuda --data-dir data --model-path pct_model.pkl --num-points 1024 --batch-size 32 --epochs 100 --seed 42

也可以使用脚本：

CUDA_VISIBLE_DEVICES=0 bash scripts/train.sh

关键参数均可通过命令行配置：

--num-points：每个点云采样点数，默认 1024
--batch-size：批大小，默认 32
--epochs：训练轮数，默认 100
--lr / --min-lr：初始学习率和最小学习率
--seed：随机种子，默认 42
--resume：从 --model-path 已有权重继续训练

训练会在 outputs/train_config.json 和 outputs/train_command.txt 保存实际配置与命令，便于复现。

推理 / 评测

使用训练得到的 pct_model.pkl 预测测试集：

CUDA_VISIBLE_DEVICES=0 cuda_archs=80 .venv/bin/python pct.py --mode test --cuda --data-dir data --model-path pct_model.pkl --output result.json --num-points 1024 --batch-size 32 --seed 42

也可以使用脚本：

CUDA_VISIBLE_DEVICES=0 bash scripts/predict.sh

若平台要求 zip 文件，可执行：

zip -j result.zip result.json

结果说明

模型输出为 40 类 logits，训练指标为交叉熵损失和训练集分类准确率。测试集提交文件 result.json 是样本编号到预测类别的字典：

{
  "0": 0,
  "1": 0,
  "2": 0
}

一次训练记录中，模型训练到第 92 个 epoch，后期训练指标约为：

epoch 91: train_acc=0.8651, train_loss=0.4149
epoch 92: train_acc=0.8654, train_loss=0.4236

使用该 checkpoint 生成 result.zip 并提交到头歌平台，评测编号为 2026042620092189213316，线上分数为 0.8416。

代码来源与许可

本仓库核心代码为 PA3 作业实现，模型结构参考 Point Cloud Transformer 的常见实现方式，使用 Jittor API 重写。项目以 MIT License 开源。