PCT_jittor

基于计图（Jittor）框架的 Point Cloud Transformer 实现，用于 ModelNet40 三维形状分类。

项目简介

本项目使用 Jittor 深度学习框架，训练 PCT（Point Cloud Transformer）模型在 ModelNet40 数据集上完成 40 类三维形状分类任务。

• 模型：Point Cloud Transformer，采用 offset attention 机制
• 数据集：ModelNet40（训练集 9,843 个，测试集 2,468 个，每个样本 2,048 个点）
• 准确率：验证集 88%+

环境配置

• Python 3.10+
• Jittor 1.3.11
• NumPy

pip install jittor numpy

数据集

下载 ModelNet40 点云数据并解压至 data/ 目录：

data/
├── train_points.npy   # 训练集点云 (9843, 2048, 3)
├── train_labels.npy   # 训练集标签 (9843,)
├── test_points.npy    # 测试集点云 (2468, 2048, 3)
└── categories.txt     # 40 个类别名称

使用方法

训练

python pct.py --data_dir ./data

主要参数：

参数	默认值	说明
--data_dir	./data	数据集路径
--epochs	200	训练轮数
--batch_size	64	批次大小
--lr	0.001	学习率
--seed	42	随机种子

生成预测结果

训练完成后会自动保存最佳模型并生成测试集预测结果 result.json。

相对 Baseline 的优化

方面	Baseline	本实现
优化器	SGD (lr=0.01)	AdamW (lr=0.001)
学习率调度	余弦退火	预热 + 余弦退火
数据增强	仅 Y 轴旋转	Y 轴旋转
验证集	无	10% 留出验证
模型选择	最后一个 epoch	验证集最佳 epoch

项目结构

.
├── pct.py             # 训练主脚本
├── precompute_aug.py  # 可选：预生成增强数据
├── README.md
├── .gitignore
└── data/              # 数据集（需单独下载，不入库）