PCT: Point Cloud Transformer for ModelNet40 Classification

项目简介

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

PCT 是一个融合了自注意力机制、几何感知和通道注意力的点云深度学习模型，能够有效地学习点云的全局和局部几何特征，在三维形状分类中取得优异性能。

文件结构

warm-up_round/
├── pct.py                  # 改进完整版本
├── README.md               # 本文件
└── data/
    ├── train_points.npy    # 训练点云数据 (N, 2048, 3)
    ├── train_labels.npy    # 训练标签 (N,)
    ├── test_points.npy     # 测试点云数据 (M, 2048, 3)
    └── categories.txt      # 类别名称

pct.py 相对于 pct_ori.py 的主要改进

1. 完整的数据增强策略

✨ 改进内容：

✅ 点云中心化（减去质心）
✅ 半径标准化（按最大距离缩放）
✅ 随机旋转（绕 Y 轴 0~2π 之间的任意角度）
✅ 随机缩放（0.9~1.1 倍）
✅ 随机抖动（±0.02 范围内的高斯噪声）

效果：提高模型的鲁棒性和泛化能力

2. 增强的模型架构

原版架构

Conv1d(3→128) → BatchNorm → ReLU
    ↓
SA_Layer × 4 (自注意力层堆叠)
    ↓
Conv1d(512→1024) → 全局最大池化
    ↓
FC layers → 分类输出

改进版架构

Conv1d(3→128) + 位置嵌入(pos_embed) → 预精化(pre_refine)
    ↓
GeometryAwareAttention × 2 (几何感知自注意力)
    ↓
SA_Layer × 2 (标准自注意力)
    ↓
注意力精化(attn_refine) → 通道注意力(ChannelAttention1D)
    ↓
全局特征融合（最大池化 + 平均池化）
    ↓
FC layers with Dropout → 分类输出

新增关键模块

a) 几何感知注意力 (GeometryAwareAttention)

在自注意力基础上加入坐标偏差项，使模型能更好地感知点云的空间几何结构
参数化：自注意力 + 坐标感知偏差

b) 通道注意力 (ChannelAttention1D)

SE（Squeeze-and-Excitation）风格的通道注意力
动态重加权各个特征通道，强化重要特征

c) 残差点 MLP (ResidualPointMLP)

逐点特征精化的残差连接模块
提升特征表达能力

d) 位置嵌入 (pos_embed)

显式地编码点云坐标信息
帮助模型学习更好的几何先验

3. 学习率调度优化

改进项：

✅ 添加 Warmup 阶段（默认10个epoch）
- 从 base_lr / warmup_epochs 逐步升到 base_lr
- 避免训练初期梯度爆炸，提高训练稳定性
✅ 更精细的余弦退火调度
- Warmup 后使用余弦衰减，平滑降低学习率

# 原版：直接余弦退火
lr = eta_min + (base_lr - eta_min) * (1 + cos(π * epoch / T_max)) / 2

# 改进版：Warmup + 余弦退火
if epoch < warmup_epochs:
    lr = base_lr * (epoch + 1) / warmup_epochs
else:
    lr = eta_min + (base_lr - eta_min) * (1 + cos(π * (epoch - warmup) / (T_max - warmup))) / 2

4. 损失函数升级

标签平滑交叉熵损失 (Label Smoothing)

✅ 结合标准交叉熵和均匀分布的KL散度
✅ 防止模型过度自信，提高泛化性能
✅ 可配置平滑系数（默认0.1）

loss = (1 - smoothing) * ce_loss + smoothing * kl_loss

5. 测试时增强 (Test-Time Augmentation)

投票平均预测 (Vote Averaging)

✅ 支持多轮旋转预测（默认10次）
✅ 绕 Y 轴等间隔旋转，每轮预测概率求和
✅ 最终取概率最高的类别，显著提升测试准确率

效果：同一样本多角度预测 → 更鲁棒的结果

6. 优化器和超参数灵活性

新增选项：

✅ 支持 SGD 和 Adam 两种优化器
✅ 可配置的 momentum（动量）和 weight_decay（权重衰减）
✅ 更丰富的命令行参数

# 示例：使用 Adam 优化器
python pct.py --optimizer adam --lr 0.001 --weight_decay 1e-4

核心参数对比

特性	pct_ori.py	pct.py
数据增强	仅旋转	旋转 + 缩放 + 抖动 + 标准化
注意力模块	SA_Layer × 4	GeometryAwareAttention × 2 + SA_Layer × 2
特征融合	仅最大池化	最大 + 平均池化
通道注意力	❌	✅
位置嵌入	❌	✅
预/后处理	❌	pre_refine + attn_refine
损失函数	交叉熵	标签平滑交叉熵
Warmup	❌	✅
投票预测	单次	多轮投票平均
优化器	SGD 仅	SGD + Adam

使用方法

安装依赖

pip install jittor numpy

运行训练

# 基础训练（使用默认参数）
python pct.py

# 自定义参数示例
python pct.py \
  --data_dir ./data \
  --n_points 1024 \
  --batch_size 32 \
  --epochs 200 \
  --lr 0.001 \
  --optimizer adam \
  --warmup_epochs 10 \
  --label_smoothing 0.1 \
  --vote_num 10 \
  --seed 42

命令行参数

参数	默认值	说明
`--data_dir`	`./data`	数据集路径
`--n_points`	`1024`	采样点数
`--batch_size`	`32`	批量大小
`--epochs`	`200`	训练轮数
`--lr`	`0.001`	学习率
`--optimizer`	`adam`	优化器 (`adam` 或 `sgd`)
`--weight_decay`	`1e-4`	权重衰减
`--momentum`	`0.9`	SGD 动量
`--eta_min`	`1e-5`	最小学习率
`--warmup_epochs`	`10`	Warmup 轮数
`--label_smoothing`	`0.1`	标签平滑系数
`--vote_num`	`10`	投票轮数
`--no_tta_rotate`	-	禁用测试时旋转增强
`--seed`	`42`	随机种子

输出结果

训练完成后生成：

pct_model.pkl - 训练好的模型权重
result.json - 测试集预测结果，格式为 {"sample_id": class_id, ...}

性能提升总结

改进版（pct.py）相对原版（pct_ori.py）的优势：

方面	提升策略	预期效果
训练稳定性	Warmup + 余弦退火	收敛更快、更稳定
模型容量	几何感知注意力 + 通道注意力	更强的特征学习能力
数据利用	完整的数据增强	更好的泛化性能
预测鲁棒性	投票平均 (TTA)	测试准确率显著提升
正则化	标签平滑 + Dropout	减少过拟合
灵活性	多种优化器和超参数	支持更多实验配置

依赖

Python >= 3.7
Jittor - 计图深度学习框架
NumPy - 数值计算库

参考

PCT: Point Cloud Transformer
ModelNet40: A 3D Object Recognition Benchmark

备注

建议使用 GPU 加速训练（Jittor 自动配置 CUDA）
数据点云已预处理为 2048 个点，采样时随机选择 1024 个点
类别总数为 40，对应 ModelNet40 数据集