PCT_jittor

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

本项目是基于 Jittor 框架实现的 Point Cloud Transformer (PCT) 点云分类 baseline，用于 ModelNet40 三维形状分类任务。仓库按照计图比赛开源规范组织，包含配置文件、核心代码、运行脚本、数据说明和可复现输出约定。

环境安装

Python: 3.8+
Framework: Jittor

安装依赖：

pip install -r requirements.txt

如需 GPU 训练，请先确认 CUDA 环境可被 Jittor 正确识别。Jittor 安装说明可参考官方文档：https://cg.cs.tsinghua.edu.cn/jittor/download/

数据准备

本项目默认读取预处理后的 NumPy 点云文件。请在仓库根目录创建 data/，并放入：

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

文件说明：

文件	说明
`train_points.npy`	训练点云，shape 通常为 `(N, 2048, 3)`
`train_labels.npy`	训练标签，shape 通常为 `(N,)`
`test_points.npy`	测试点云，shape 通常为 `(N, 2048, 3)`

数据根目录通过配置文件字段 data_dir 设置，也可以用命令行参数覆盖：

python src/pct.py --config configs/default.json --data_dir ./data

训练

使用默认配置训练：

python src/pct.py --config configs/default.json --mode train

覆盖部分配置示例：

python src/pct.py \
  --config configs/default.json \
  --mode train \
  --data_dir ./data \
  --output_dir ./outputs \
  --epochs 200 \
  --batch_size 32 \
  --lr 0.01 \
  --seed 42

也可以使用脚本：

bash scripts/train.sh

评测/推理

训练会在 outputs/ 中保存模型并自动生成一次测试集预测结果。单独推理可执行：

python src/pct.py \
  --config configs/default.json \
  --mode infer \
  --ckpt outputs/pct_model.pkl \
  --result_file outputs/result.json

也可以使用脚本：

bash scripts/infer.sh

结果说明

训练日志会输出训练集分类准确率：

Train Acc = 正确分类样本数 / 训练样本总数 * 100%

推理结果保存为 outputs/result.json，格式为：

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

其中 key 为测试样本编号，value 为预测类别编号。

可复现设置

项目提供 --seed 参数，并在代码中统一设置 NumPy 与 Jittor 随机种子。每次运行会在 output_dir 中保存：

config.json：实际使用的配置
command.txt：实际运行命令
train.log：训练日志，仅训练模式生成
pct_model.pkl：模型权重，仅训练模式生成
result.json：测试集预测结果

命令行参数优先级高于 configs/default.json。

项目结构

PCT_jittor/
├── configs/
│   └── default.json
├── data/
│   └── README.md
├── scripts/
│   ├── train.sh
│   └── infer.sh
├── src/
│   └── pct.py
├── tools/
│   └── README.md
├── .gitignore
├── LICENSE
├── README.md
└── requirements.txt

代码说明

src/pct.py 中主要包含：

ModelNet40Dataset：读取 ModelNet40 点云数据，支持训练数据增强
SA_Layer：PCT self-attention 层
PCT：点云分类模型主体
CosineAnnealingLR：余弦退火学习率调度器
train_one_epoch：单轮训练逻辑
predict：测试集推理并导出预测结果

第三方引用与声明

本项目代码用于课程作业与 Jittor 点云分类实验学习。PCT 模型思路来源于 Point Cloud Transformer 相关工作，主体代码框架来源于清华大学计算机图形学基础课程。