PCT_jittor

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

基于 Jittor 实现的 Point Cloud Transformer（PCT），在 ModelNet40 点云数据集上完成 40 类三维形状分类。适用于清华大学计算机图形学 PA3 / 计图挑战赛赛道二热身赛。

功能简介

加载预处理好的 ModelNet40 点云（.npy）
PCT 分类网络（Conv1d 嵌入 + 4 层 Self-Attention + 全局池化 + FC）
训练阶段数据增强（Y 轴旋转、小角度扰动、随机缩放、高斯抖动）
SGD 优化器 + 余弦退火学习率
训练后导出 pct_model.pkl 与测试集预测 result.json

环境要求

Linux / WSL（推荐）
Python 3.7+
C++编译器（g++）
CUDA（可选，用于 GPU 加速）
Jittor

pip install jittor numpy
# 或使用 conda
conda create -n jittor python=3.7
conda activate jittor
pip install jittor

数据集

本仓库不包含大型 .npy 数据文件。请从课程/比赛页面下载 ModelNet40 点云数据，解压到项目根目录下的 data/ 文件夹，目录结构如下：

data/
├── train_points.npy   # (9843, 2048, 3)
├── train_labels.npy   # (9843,)
├── test_points.npy    # (2468, 2048, 3)
└── categories.txt     # 40 类名称（已随仓库提供）

下载入口（以课程页面为准）：

计图挑战赛示例代码附件 / Educoder 竞赛页

快速开始

git clone https://gitlink.org.cn/LuckyBoy/PCT_jittor_lry24.git
cd PCT_jittor_lry24
conda activate jittor   # 或你的 jittor 环境

# 确认数据已就位
ls data/train_points.npy data/train_labels.npy data/test_points.npy

# 训练 + 测试集预测（默认 200 epoch）
python pct.py

# 可选：将日志同时保存到文件
python pct.py 2>&1 | tee train.log

训练结束后将生成：

文件	说明
`pct_model.pkl`	模型权重
`result.json`	测试集预测，格式 `{"0": 4, "1": 35, ...}`

打包评测提交：

zip -j result.zip result.json

命令行参数

python pct.py --help

参数	默认值	说明
`--data_dir`	`./data`	数据目录
`--n_points`	`1024`	每样本采样点数
`--batch_size`	`32`	批大小
`--epochs`	`200`	训练轮数
`--lr`	`0.01`	初始学习率
`--seed`	`42`	随机种子

示例：

python pct.py --epochs 1              # 快速冒烟测试
python pct.py --n_points 2048         # 使用更多点（更慢、可能更准）
python pct.py --batch_size 16         # 显存不足时减小 batch

项目结构

PCT_jittor/
├── pct.py              # 数据加载、PCT 模型、训练与推理
├── README.md           # 本文件
├── .gitignore
├── data/
│   └── categories.txt  # 类别名（数据 .npy 需自行下载）
└── docs/
    ├── README.md       # 赛题说明（热身赛）
    └── PA3_README.pdf  # 课程讲义（可选）

模型说明

输入点云形状为 (B, 3, N)，经两层 Conv1d 与四个 SA_Layer 自注意力模块后，拼接多层特征并做全局 max pooling，最后经全连接输出 40 维 logits。优化器为 SGD（momentum=0.9, weight_decay=1e-4），学习率按 余弦退火 从 lr 衰减至 1e-4。

参考论文：PCT: Point Cloud Transformer

性能参考

在完整训练（200 epoch，1024 点，默认增强与超参）下，训练集准确率通常可达 90% 以上；计图热身赛测试集准确率 ≥ 80% 为通过线（具体以评测平台为准）。

许可证

本项目为课程作业开源，仅供学习交流。使用 Jittor 请遵循其官方许可与引用说明。

引用

若参考 PCT 原论文，可引用：

@article{Guo_2021,
  title={PCT: Point cloud transformer},
  journal={Computational Visual Media},
  volume={7},
  number={2},
  pages={187--199},
  year={2021}
}