PCT Jittor

A Jittor Implementation of Point Cloud Transformer for 3D Shape Classification

85.01% Test Accuracy · ModelNet40 · 200 Epochs · ~28s/epoch

Quick Start

python pct.py --epochs 200 --batch_size 32 --lr 0.01

_{训练完成后自动生成 pct_model.pkl 和 result.json}

简介

本项目基于 Jittor 深度学习框架，复现了论文 PCT: Point Cloud Transformer (CVMJ 2021) 中提出的三维点云分类网络，并在 ModelNet40 数据集上完成训练与评测。

PCT 将 Transformer 的自注意力机制引入点云处理，通过全局关系建模替代传统的局部邻域聚合，在三维形状分类任务上取得了优异的性能。

网络架构

PCT 网络架构示意图

Self-Attention Layer 内部结构：

步骤	操作	输出维度
Q/K 映射（共享权重）	Conv1d(C → C//4)	(B, C//4, N)
V 映射	Conv1d(C → C)	(B, C, N)
注意力矩阵	Softmax(Q^T · K)	(B, N, N)
特征聚合	V · A	(B, C, N)
残差输出	F + ReLU(BN(W_T(F − F_attn)))	(B, C, N)

环境配置

依赖：

Python 3.9+
Jittor >= 1.3.11
NumPy
CUDA 12.x（推荐，亦支持 CPU）

安装 Jittor：

pip install jittor

# 验证 CUDA 是否可用
python -m jittor.test.test_cudnn_op

详细安装说明：https://cg.cs.tsinghua.edu.cn/jittor/download/

数据集

ModelNet40 数据集（npy 格式）从头歌平台下载：

https://www.educoder.net/competitions/Jittor-7

下载后放置于 data/ 目录：

PCT_jittor/
├── pct.py
├── data/
│   ├── train_points.npy    # (9843, 2048, 3)  训练集点云
│   ├── train_labels.npy    # (9843,)           训练集标签
│   └── test_points.npy     # (2468, 2048, 3)  测试集点云

每个点云原始包含 2048 个点，训练时随机采样 1024 个点作为网络输入。

使用方法

训练

python pct.py --epochs 200 --batch_size 32 --lr 0.01

后台运行（推荐）：

nohup python pct.py --epochs 200 --batch_size 32 --lr 0.01 > train.log 2>&1 &
tail -f train.log

参数说明

参数	默认值	说明
`--epochs`	200	训练轮数
`--batch_size`	32	批大小
`--lr`	0.01	初始学习率（余弦退火至 1e-3）
`--n_points`	1024	每个点云的采样点数
`--data_dir`	`./data`	数据集目录
`--seed`	42	随机种子

输出文件

文件	说明
`pct_model.pkl`	训练完成的模型权重
`result.json`	测试集预测结果，格式：`{"0": 类别, "1": 类别, ...}`

训练细节

配置项	值
优化器	SGD（momentum=0.9，weight_decay=1e-4）
学习率调度	CosineAnnealingLR（初始 0.01 → 最终 1e-3）
数据增强	随机绕 Y 轴旋转 + 高斯抖动（σ=0.01，clip ±0.05）
每 epoch 用时	~28 秒（RTX 40 系 GPU）
总训练时间	~95 分钟

训练曲线（Train Accuracy）：

Epoch   1 :  12%
Epoch  50 :  65%
Epoch 100 :  80%
Epoch 150 :  86%
Epoch 200 :  89%  ← 最终

结果

指标	数值
Train Accuracy（Epoch 200）	89.04%
Test Accuracy（头歌平台评测）	85.01%
模型参数量	1.37M
测试集样本数	2468

项目结构

PCT_jittor/
├── pct.py          # 完整实现：数据集 / 模型 / 训练 / 预测
├── README.md
├── .gitignore
├── data/           # 数据集（不上传，请自行下载）
│   ├── train_points.npy
│   ├── train_labels.npy
│   └── test_points.npy
├── pct_model.pkl   # 训练权重（不上传）
└── result.json     # 预测结果（不上传）

参考

@article{guo2021pct,
  title   = {PCT: Point Cloud Transformer},
  author  = {Guo, Meng-Hao and Cai, Jun-Xiong and Liu, Zheng-Ning and
             Mu, Tai-Jiang and Martin, Ralph R and Hu, Shi-Min},
  journal = {Computational Visual Media},
  year    = {2021}
}

论文：https://arxiv.org/abs/2012.09688
Jittor 框架：https://cg.cs.tsinghua.edu.cn/jittor/
ModelNet40 评测平台：https://www.educoder.net/competitions/Jittor-7