PCT (Point Cloud Transformer) in Jittor for ModelNet40 Classification

项目简介

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

模型通过 Self-Attention 机制建模点与点之间的关系，相比传统 PointNet 类方法，能够更好地捕捉全局几何结构信息。

---

模型特点

• 使用 Transformer-style Self-Attention 进行点云特征建模
• 支持 端到端训练
• 采用 多层 Attention + 特征融合 提升表达能力
• 使用 Cosine Annealing 学习率调度 提升收敛效果

---

项目结构

.
├── pct.py                 # 主程序（模型 + 训练 + 推理）
├── data/
│   ├── train_points.npy   # 训练点云数据 (N, 2048, 3)
│   ├── train_labels.npy   # 训练标签 (N,)
│   └── test_points.npy    # 测试点云数据
├── pct_model.pkl          # 训练后模型（自动生成）
└── result.json            # 测试预测结果（自动生成）

---

环境依赖

pip install jittor numpy

---

运行方法

训练模型

python pct.py

可选参数：

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

---

数据说明

输入数据格式

• train_points.npy：

  (N, 2048, 3)

• train_labels.npy：

  (N,)

• test_points.npy：

  (N, 2048, 3)

  其中的3表示xyz坐标

---

数据增强策略

训练阶段包含以下增强方法：

• 随机旋转（绕 Y 轴）
• 随机缩放
• 随机平移
• 随机抖动（Gaussian noise）

---

模型结构

输入

(B, 3, N)

核心模块

• Conv1d 特征提取
• 4 层 Self-Attention
• 特征拼接与融合
• 全局最大池化
• 全连接分类器

---

损失函数

使用：

Cross Entropy Loss

---

优化器与调度

• 优化器：SGD（momentum=0.9）
• 学习率策略：Cosine Annealing

可选优化器：Adam

替换为

optimizer = nn.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4)

---

推理结果

运行完成后会生成：

result.json

格式如下：

{
  "0": 12,
  "1": 3,
  "2": 25
}

key：样本索引 value：预测类别

---

模型保存

训练结束后自动保存：

pct_model.pkl

---

注意事项

• 输入必须为 (B, N, 3)，模型内部会转换为 (B, 3, N) 使用permute(0,2,1)
• labels 必须为一维 (B,)
• GPU 训练需开启：

jt.flags.use_cuda = 1

---

总结

本项目实现了一个完整的点云分类 pipeline，包括：

• 数据加载与增强
• Transformer 模型构建
• 训练与优化
• 测试预测与结果输出

---