PCT_jittor

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

项目简介

本项目基于计图（Jittor）深度学习框架，实现了 PCT（Point Cloud Transformer） 模型，用于 ModelNet40 三维点云形状分类任务。PCT 将 Transformer 中的自注意力（Self-Attention）机制引入点云处理领域，通过全局注意力建模点之间的长距离依赖关系，在点云分类与分割任务上取得了优异性能。

项目的代码框架来源于 头歌平台计图挑战赛：第六届计图人工智能挑战赛

参考论文： https://doi.org/10.48550/arXiv.2012.09688

在清华大学 《计算机图形学》 的课程指导下完成了本项目

环境依赖

• Python 3.7+
• Jittor >= 1.3
• NumPy

数据集

使用 ModelNet40 数据集，存放于 data/ 目录

项目结构

.
├── data/
│   ├── train_points.npy      # 训练集点云数据
│   ├── train_labels.npy      # 训练集标签
│   ├── test_points.npy       # 测试集点云数据
│   └── categories.txt        # 类别名称列表
├── pct.py                    # 训练与预测主程序
├── pct_model.pkl             # 训练完成的模型权重
├── result.json               # 测试集预测结果
├── result.zip                # 预测结果压缩包
└── README.md                 

使用方法

训练模型

python pct.py

默认超参数：

参数	默认值	说明
--data_dir	./data	数据集路径
--n_points	1024	采样点数
--batch_size	32	批大小
--epochs	200	训练轮数
--lr	0.01	初始学习率
--seed	42	随机种子

可通过命令行参数自定义配置，例如：

python pct.py --epochs 300 --lr 0.001 --batch_size 64

训练完成后，模型权重将保存为 pct_model.pkl。

仅执行预测

训练完成后程序将自动对测试集进行预测，结果保存至 result.json。如需使用已训练的模型单独预测，可修改 main() 函数跳过训练阶段。

实现细节

数据增强

训练时对每个点云随机应用以下变换（仅训练集）：

• 随机旋转：绕 Y 轴旋转 θ ∼ U(0, 2π)
• 随机缩放：缩放因子 s ∼ U(0.8, 1.25)
• 随机平移：平移量 t ∼ U(−0.1, 0.1)
• 随机抖动：高斯噪声 N(0, 0.01) 截断至 [−0.05, 0.05]

训练策略

• 优化器：SGD，momentum=0.9，weight_decay=1e−4
• 学习率调度：Cosine Annealing，从 0.01 余弦衰减至 1e−5（T_max=200）
• 损失函数：Cross Entropy Loss
• 正则化：Dropout（p=0.5）+ L2 权重衰减

参考资料

• PCT: Point Cloud Transformer (arXiv:2012.09688)
• Jittor 官方文档
• Educoder 竞赛页面