ModelNet40 Point Cloud Classification - PCT (Point Cloud Transformer)

基于 Jittor 框架的 PCT (Point Cloud Transformer) 三维点云分类模型，用于 ModelNet40 数据集的形状分类任务。

环境安装

Python 版本

Python 3.7+

安装依赖

pip install jittor numpy

注意：Jittor 需要根据您的系统环境选择合适的版本。建议参考 Jittor 官方文档进行安装。

数据准备

数据集说明

本项目使用 ModelNet40 数据集，包含 40 个类别的三维 CAD 模型。

数据目录结构

PointCloud/
├── data/
│   ├── categories.txt          # 40 个类别名称
│   ├── train_points.npy        # 训练集点云数据 (N, 2048, 3)
│   ├── train_labels.npy        # 训练集标签 (N,)
│   └── test_points.npy         # 测试集点云数据 (M, 2048, 3)
├── pct.py                      # 主程序
├── README.md                   # 本文件
├── LICENSE                     # 许可证
└── .gitignore                  # Git 忽略文件

数据格式

点云数据: .npy 格式，每个样本包含 2048 个三维点坐标 (x, y, z)
标签数据: 整数标签，范围 0-39，对应 40 个类别
类别列表: categories.txt 按行存储类别名称

训练

基本训练命令

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

参数说明

参数	类型	默认值	说明
`--data_dir`	str	`./data`	数据集根目录
`--n_points`	int	1024	每个样本采样的点数
`--batch_size`	int	32	批大小
`--epochs`	int	200	训练轮数
`--lr`	float	0.001	初始学习率
`--seed`	int	42	随机种子

训练过程

优化器: AdamW (weight_decay=1e-2)
学习率调度: 余弦退火 (Cosine Annealing) + 线性预热 (Warmup)
预热轮数: 10 epochs
数据增强: 随机旋转、随机抖动、随机缩放、随机丢弃

训练过程中会输出：

每个 batch 的损失值和准确率
每个 epoch 的训练损失、训练准确率、当前学习率和耗时
最终保存模型到 pct_model.pkl

评测/推理

生成预测结果

训练完成后，程序会自动对测试集进行预测并生成 result.json：

python pct.py --data_dir ./data --epochs 200

预测结果格式

result.json 文件格式：

{
  "0": 15,
  "1": 28,
  "2": 5,
  ...
}

Key: 样本编号 (字符串)
Value: 预测类别编号 (0-39 的整数)

模型结构

PCT (Point Cloud Transformer) 架构

输入: (B, 3, N) 点云坐标
  ↓
Conv1d(3 → 128) + BatchNorm + ReLU
  ↓
Conv1d(128 → 128) + BatchNorm + ReLU
  ↓
Self-Attention Layer × 4      # 多头自注意力机制
  ↓
Concat([x1, x2, x3, x4])      # 特征融合 (B, 512, N)
  ↓
Conv1d(512 → 1024) + BatchNorm + LeakyReLU
  ↓
Max Pooling                    # (B, 1024)
  ↓
FC(1024 → 512) + BatchNorm + ReLU + Dropout(0.5)
  ↓
FC(512 → 256) + BatchNorm + ReLU + Dropout(0.5)
  ↓
FC(256 → 40)                   # 输出类别 logits

模型参数

约 3.5M 可训练参数

结果说明

评估指标

准确率 (Accuracy): 正确预测的样本数 / 总样本数
类别: 40 个 ModelNet40 类别

预期性能

在完整训练 (200 epochs) 后，预期可以达到：

训练集准确率: 85%+ (取决于超参数和数据增强策略)
推理速度: 约 0.1s/批次 (batch_size=32)

注意事项

由于训练轮数、数据增强策略和随机种子的差异，实际结果可能与预期存在合理差异
建议使用 GPU 进行训练和推理以获得更好的性能
测试集预测结果保存在 result.json 中

可复现性

随机种子设置

代码中统一设置随机种子以确保结果可复现：

np.random.seed(args.seed)
jt.set_global_seed(args.seed)

配置保存

建议在训练时记录以下信息：

实际使用的命令行参数
训练日志 (自动输出到控制台)
模型权重文件 (pct_model.pkl)

项目结构说明

核心文件

pct.py: 主程序文件，包含数据加载、模型定义、训练和推理逻辑
data/: 数据目录 (包含预处理后的 .npy 文件)
result.json: 测试集预测结果 (训练后自动生成)

代码特点

使用 Jittor 框架实现
模块化设计，易于扩展
包含完整的数据增强策略
支持命令行参数配置
统一的随机种子设置

许可证

本项目采用 MIT 许可证，详见 LICENSE 文件。

引用

本项目参考了以下工作：

PCT: Point Cloud Transformer [arXiv:2012.09688]
ModelNet40: 3D Shape Classification Dataset [arXiv:1512.03012]

贡献

欢迎提交 Issue 和 Pull Request 来改进本项目。

联系方式

如有问题或建议，请通过 GitHub Issue 联系我们。