PCT_jittor_zhangzx

A Jittor implementation of Point Cloud Transformer (PCT) for ModelNet40 classification

---

项目简介

本项目实现了Point Cloud Transformer (PCT)网络，用于三维点云的形状分类任务。PCT是一种基于自注意力机制的点云处理网络，能够有效建模点云中点与点之间的全局依赖关系。

主要特点：

• 纯Transformer架构，完全基于自注意力机制
• 多层特征融合，捕捉多尺度语义信息
• 基于Jittor框架，支持GPU加速和混合精度训练
• 集成随机旋转、缩放、噪声等数据增强策略

---

模型架构

Input (B, 3, N)
↓
Conv1d + BN + ReLU (3 → 128)
↓
Conv1d + BN + ReLU (128 → 128)
↓
SA_Layer × 4 (Self-Attention)
↓
Concatenate (512, N)
↓
Conv1d + BN + LeakyReLU (512 → 1024)
↓
Max Pooling (global)
↓
Linear × 3 (1024 → 512 → 256 → 40)
↓
Output (B, 40) logits

Self-Attention Layer结构：

Input (B, C, N)
↓
Q = Conv_q(F)  (C → C//4)
K = Conv_k(F)  (C → C//4)
V = Conv_v(F)  (C → C)
↓
Attention = Softmax(Q^T · K)
↓
F_attn = V · Attention
↓
Output = F + MLP(F - F_attn)

---

快速开始

环境要求

• Python 3.8+
• Jittor >= 1.3.0
• NumPy
• 可选：CUDA（GPU加速）

安装依赖

pip install jittor

# 验证安装
python -c "import jittor as jt; print(jt.flags.use_cuda)"

数据准备

数据集结构：

data/
├── train_points.npy    # 训练集点云 (N_train, 2048, 3)
├── train_labels.npy    # 训练集标签 (N_train,)
└── test_points.npy     # 测试集点云 (N_test, 2048, 3)

训练模型

# 基础训练
python pct.py

# 自定义参数
python pct.py --epochs 200 --batch_size 64 --lr 0.001 --n_points 1024

# 使用GPU
python pct.py --batch_size 128

命令行参数

参数 类型 默认值 描述 –data_dir str ./data 数据目录路径 –n_points int 1024 采样点数 –batch_size int 32 批次大小 –epochs int 200 训练轮数 –lr float 0.01 初始学习率 –seed int 42 随机种子

---

训练输出示例

============================================================
ModelNet40 Classification - PCT Baseline
Points: 1024  Batch: 64  Epochs: 200  LR: 0.01
============================================================
Train: 9843 samples
Test:  2468 samples
Model parameters: 1.37M

  Epoch [1] Batch [20] Loss: 3.5621  Acc: 9.38%  Time: 12.3s
  Epoch [1] Batch [40] Loss: 3.4012  Acc: 13.28%  Time: 24.1s
...
Epoch [200/200] Loss: 0.4521  Train Acc: 89.23%  LR: 0.000010  Time: 35.2s

预期性能（GPU环境）

GPU 每个Epoch 200 Epochs 期望准确率 RTX 3060 约35秒 约2小时 88-90% RTX 3080 约20秒 约1.1小时 88-90% V100 约10秒 约35分钟 89-91%

---

输出文件

训练完成后生成：

文件 描述 pct_model.pkl 训练好的模型权重 result.json 测试集预测结果

result.json格式：

{
  "0": 12,
  "1": 5,
  "2": 28,
  ...
  "2467": 3
}

---

性能优化

启用混合精度训练

在main()函数开头添加：

jt.flags.amp_level = 1

效果：显存占用减少40%，训练速度提升30-50%

优化数据加载

train_loader = ModelNet40Dataset(
    # ...
    num_workers=8,
)

调整批次大小

python pct.py --batch_size 128

---

数据增强

训练时自动应用的增强策略：

增强类型 参数范围 概率 绕Y轴旋转 02π 100% 绕X轴旋转 -π/4 ~ π/4 50% 绕Z轴旋转 -π/4 ~ π/4 50% 随机缩放 0.81.2 100% 高斯噪声 σ=0.02 50%

---

故障排除

编译卡在 “Compiling Operators…”

原因：Jittor首次运行时的即时编译，正常现象

解决：耐心等待1-2分钟，编译完成后会缓存结果

GPU不可用

python -c "import jittor as jt; print(jt.flags.use_cuda)"

如果返回False，检查Jittor安装或CUDA环境

显存不足（OOM）

# 减小batch_size
python pct.py --batch_size 16

# 或减少采样点数
python pct.py --n_points 512

---

项目结构

PCT_Jittor/
├── pct.py              # 主程序：模型定义、训练、推理
├── data/               # 数据集目录
│   ├── train_points.npy
│   ├── train_labels.npy
│   └── test_points.npy
├── pct_model.pkl       # 训练好的模型（生成）
├── result.json         # 预测结果（生成）
└── README.md

---

许可证

MIT License

---

致谢

· Jittor框架 - 清华大学开源深度学习框架 · ModelNet40 - 3D点云分类基准数据集

```