PCT (Point Cloud Transformer) for ModelNet40 Classification

概述

这是一个基于 Jittor 框架的点云分类模型实现，用于 ModelNet40 三维形状分类任务。模型采用 Point Cloud Transformer (PCT) 架构，利用自注意力机制处理点云数据。

功能特性

• 点云数据集加载：支持从 .npy 文件加载预处理的点云数据
• 数据增强：包括随机旋转、缩放、平移、抖动与打乱排序
• PCT 模型：基于自注意力层的深度学习模型
• 灵活的优化器选择：支持 SGD 和 Adam 优化器，经测试后选择 SGD 优化器
• 余弦退火学习率调度：自动调整学习率以提升训练效果
• 批量预测：对测试集进行推理并输出 JSON 格式的预测结果

依赖安装

sudo apt install python3.7-dev libomp-dev
python3.7 -m pip install jittor
python3.7 -m jittor.test.test_example
# 如果您电脑包含Nvidia显卡，检查cudnn加速库
python3.7 -m jittor.test.test_cudnn_op

文件结构

PA3/
├── pct.py              # 主程序
├── README.md           # 本文件
├── data/               # 数据目录
|   ├── catagories.txt      # 物品类别标签
│   ├── train_points.npy    # 训练集点云 (N_train, 2048, 3)
│   ├── train_labels.npy    # 训练集标签 (N_train,)
│   ├── test_points.npy     # 测试集点云 (N_test, 2048, 3)
├── pct_model.pkl       # 保存的训练完成的模型（模型训练后的输出文件）
└── result.json         # 测试集预测结果（模型训练后的输出文件）

快速开始

基本使用

python pct.py

自定义参数

参数	默认值	说明
--data_dir	./data	数据集所在目录
--n_points	1024	每个点云采样的点数
--batch_size	32	批次大小
--epochs	200	训练轮数
--lr	0.01	初始学习率
--seed	42	随机种子

• 使用示例：

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

• 根据 ModelNet40 数据集特性，推荐的超参数配置：

  • 学习率：0.01 ~ 0.001（SGD）或 0.001 ~ 0.0001（Adam）
  • 动量：0.9（SGD）
  • 权重衰减：1e-4
  • Dropout：0.5
  • 批次大小：32 ~ 64
  • 训练轮数：200 ~ 300

核心组件

1. ModelNet40Dataset

点云数据集类，支持：

• 从 .npy 文件加载点云和标签
• 动态点采样（若点数不足时重复采样）
• 数据增强策略：
  • 随机绕 Y 轴旋转
  • 随机均匀缩放 (0.8 ~ 1.2x)
  • 随机平移 (±0.1)
  • 高斯噪声抖动 (±0.05)
  • 随机打乱点的顺序

2. SA_Layer (Self-Attention Layer)

自注意力层实现：

• 计算点之间的相似性（Query-Key）
• 使用 Softmax 得到注意力权重
• 通过权重聚合特征（Value）
• 残差连接和批归一化

3. PCT (Point Cloud Transformer)

完整的分类模型：

• 编码器：2 个卷积层 + 4 个自注意力层
• 融合层：多尺度特征拼接和压缩
• 分类头：全连接层 + 批归一化 + dropout

输入 (B, 3, N)
  ↓ Conv1d + BN + ReLU (× 2)
  ↓ Self-Attention (× 4)
  ↓ 拼接多尺度特征
  ↓ Conv1d 融合 (3×128 → 1024)
  ↓ 全局最大池化
  ↓ FC + BN + ReLU + Dropout
  ↓ FC + BN + ReLU + Dropout
  ↓ FC (输出: num_classes)
输出 (B, 40)

4. CosineAnnealingLR

余弦退火学习率调度器：

• 平滑地从初始学习率递减到最小值
• 公式：lr = η_min + (lr_0 - η_min) × (1 + cos(πt/T_max)) / 2

训练流程

1. 加载数据：从 data/ 目录加载点云和标签
2. 构建模型：初始化 PCT 模型（约 2M 参数）
3. 设置优化器：SGD + 余弦退火学习率调度
4. 训练：
   • 逐轮迭代数据
   • 计算交叉熵损失
   • 反向传播更新权重
   • 定期打印训练损失和准确率
5. 保存模型：训练完成后保存到 pct_model.pkl
6. 预测：在测试集上进行推理，生成 result.json

输出文件

1. pct_model.pkl

训练完成后保存的模型文件：

• 包含所有学习到的权重和网络结构
• 可用于加载已训练的模型进行推理或继续训练

2. result.json

包含测试集预测结果的 JSON 文件，格式如下：

{
  "0": 10,
  "1": 24,
  "2": 5,
  ...
  "sample_id": predicted_class
}

其中 sample_id 是测试样本的索引，predicted_class 是预测的类别编号 (0-39)。

参考

• Jittor 文档：https://github.com/Jittor/jittor
• 代码框架：第六届计图人工智能挑战赛-热身赛赛道二https://www.educoder.net/competitions/Jittor-7

许可证

仅供学习和研究使用。