PCT: Point Cloud Transformer for ModelNet40 Classification

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

项目简介

本项目实现了 PCT 模型对 ModelNet40 数据集的 40 类三维形状分类。PCT 是一种基于自注意力机制的点云处理网络，通过多层 Self-Attention 捕获点云中点与点之间的全局依赖关系，从而实现高效的形状特征提取与分类。

核心思路

1. 输入嵌入：将原始 3D 坐标通过 1D 卷积映射到 128 维特征空间
2. 自注意力编码：4 层 SA_Layer 逐层提取点云全局特征，每层使用共享权重的 Q/K 投影与归一化注意力机制
3. 特征融合：拼接 4 层注意力输出（512 维），经卷积融合至 1024 维
4. 全局池化：最大池化得到全局特征向量
5. 分类头：全连接层 + BatchNorm + Dropout 输出 40 类预测

模型架构

Input (B, 3, 1024)
    │
    ├── Conv1d(3→128) + BN + ReLU
    ├── Conv1d(128→128) + BN + ReLU
    │
    ├── SA_Layer(128) → x1
    ├── SA_Layer(128) → x2
    ├── SA_Layer(128) → x3
    ├── SA_Layer(128) → x4
    │
    ├── Concat(x1,x2,x3,x4) → (B, 512, N)
    ├── Conv1d(512→1024) + BN + LeakyReLU
    ├── MaxPool → (B, 1024)
    │
    ├── Linear(1024→512) + BN + ReLU + Dropout(0.5)
    ├── Linear(512→256)  + BN + ReLU + Dropout(0.5)
    └── Linear(256→40)   → logits

SA_Layer（自注意力层）

• Q/K 共享权重投影至 C/4 维，V 投影至 C 维
• Softmax 注意力 + 行归一化
• 残差连接：x = x + ReLU(BN(Conv(x - Attention(x))))

项目结构

PA3/
├── data/
│   ├── categories.txt      # 40 个类别名称
│   ├── train_points.npy    # 训练集点云 (9843, 2048, 3)
│   ├── train_labels.npy    # 训练集标签 (9843,)
│   └── test_points.npy     # 测试集点云 (2468, 2048, 3)
├── pct.py                  # 主程序（模型 + 训练 + 推理）
├── pct_best.pkl            # 最佳模型权重
├── pct_model.pkl           # 最终模型权重
├── result.json             # 测试集预测结果
└── README.md               # 项目说明文档

快速开始

环境依赖

• Python >= 3.8
• Jittor >= 1.3
• NumPy
• CUDA

安装

pip install jittor numpy

训练与预测

# 使用默认参数训练（250 epochs, SGD + Cosine Annealing）
python pct.py

# 自定义参数
python pct.py --epochs 250 --batch_size 32 --lr 0.01 --n_points 1024

训练完成后，程序将自动：

1. 保存最佳模型到 pct_best.pkl
2. 保存最终模型到 pct_model.pkl
3. 加载最佳模型，对测试集进行 10 次 softmax 概率投票
4. 输出预测结果到 result.json

命令行参数

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

训练策略

优化器与学习率

• 优化器：SGD（momentum=0.9, weight_decay=1e-4）
• Warmup：前 10 个 epoch 线性增长学习率（0 → 0.01）
• 调度器：Cosine Annealing（T_max=240, eta_min=1e-4）

数据增强

训练时对每个点云样本施加以下增强：

增强方式	参数	说明
随机旋转	Y 轴, [0, 2π]	绕 Y 轴随机旋转
随机缩放	[0.8, 1.2]	均匀随机缩放
随机平移	±0.1	三轴随机平移
高斯抖动	σ=0.02	添加高斯噪声

预测策略

• 加载训练过程中最佳模型（按训练准确率选择）
• 对测试集进行 10 次前向推理
• 累积 softmax 概率，取 argmax 作为最终预测
• 相比硬投票，概率投票更稳定，能更好地处理边界样本

模型参数

• 总参数量：~1.37M
• 输入：(B, 3, 1024) 点云
• 输出：(B, 40) 分类 logits

输出格式

result.json 格式为 {样本编号: 预测类别}，类别为 0-39 的整数，对应 categories.txt 中的 40 个类别：

{
  "0": 0,
  "1": 5,
  "2": 12,
  ...
}