PCT-Jittor: Point Cloud Transformer for ModelNet40 Classification

基于 Jittor 框架实现的 PCT (Point Cloud Transformer) 模型，用于 ModelNet40 三维形状分类任务。

项目结构

PCT_Jittor/
├── pct.py              # 主代码：模型定义、数据集、训练与推理
├── data/
│   ├── categories.txt  # 40 个物体类别名称
│   ├── train_points.npy
│   ├── train_labels.npy
│   └── test_points.npy
├── checkpoints/        # 模型检查点
├── result.json         # 测试集预测结果
├── LICENSE
└── README.md

项目依赖

• Python >= 3.7
• Jittor
• NumPy

安装依赖：

pip install jittor

使用说明

训练

python pct.py \
    --data_dir ./data \
    --n_points 1024 \
    --batch_size 32 \
    --epochs 200 \
    --lr 0.01 \
    --augment \
    --save_dir ./checkpoints

仅预测

加载已有 checkpoint 对测试集进行推理：

python pct.py --predict_only --ckpt ./checkpoints/best.pkl

预测结果保存至 result.json，格式为 {样本ID: 类别ID}。

主要参数

参数	默认值	说明
--data_dir	./data	数据目录
--n_points	1024	每个点云采样的点数
--batch_size	32	批大小
--epochs	200	训练轮数
--lr	0.01	初始学习率
--augment	False	启用数据增强
--save_dir	./checkpoints	检查点保存目录
--save_freq	0	每 N 轮保存一次检查点
--seed	42	随机种子

模型架构

Input (B, 3, N)
    ↓
Conv1d(3→128) + BN + ReLU
    ↓
Conv1d(128→128) + BN + ReLU
    ↓
SA_Layer × 4  →  concat  →  (B, 512, N)
    ↓
Conv1d(512→1024) + BN + LeakyReLU
    ↓
MaxPool  →  (B, 1024)
    ↓
FC(1024→512) + BN + ReLU + Dropout
    ↓
FC(512→256) + BN + ReLU + Dropout
    ↓
FC(256→40)  →  Output logits

数据增强

训练时可选的增强策略（概率触发）：

增强方式	概率	说明
Random Rotate Y	100%	绕 Y 轴随机旋转
Random Scale	20%	缩放因子 0.9 ~ 1.1
Random Jitter	20%	高斯噪声 (σ=0.01)
Random Translate	20%	平移 ±0.1