PCT_jittor

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

项目简介

本项目基于 Jittor 框架实现了 Point Cloud Transformer (PCT) 模型，用于 ModelNet40 三维点云分类任务。ModelNet40 包含 40 类三维物体（如飞机、椅子、桌子等），训练集 9843 个样本，测试集 2468 个样本，每个样本为 2048 个点的三维坐标。

PCT 通过 Self-Attention 机制捕捉点云中的全局上下文信息，相比传统基于 PointNet 的方法具有更强的特征表达能力。

环境要求

• Python 3.7+
• Jittor 深度学习框架
• Linux / Windows(WSL) / macOS

安装依赖

pip install jittor

数据集

数据集放置在 data/ 目录下：

文件	形状	说明
train_points.npy	(9843, 2048, 3)	训练集点云坐标
train_labels.npy	(9843,)	训练集标签 (0-39)
test_points.npy	(2468, 2048, 3)	测试集点云坐标
categories.txt	40 行	类别名称列表

数据已归一化至单位球，无需额外预处理。

注意：训练集和测试集的 .npy 文件较大（~289MB），未包含在仓库中。 请从 ModelNet40 下载，并将 train_points.npy、train_labels.npy、test_points.npy 放入 data/ 目录。

运行

直接运行（使用默认参数）

python pct.py

默认参数：200 轮训练，batch size 32，采样 1024 个点，Adam 优化器初始学习率 0.001。

自定义参数

python pct.py --epochs 200 --batch_size 32 --n_points 1024 --seed 42

完整参数列表：

参数	默认值	说明
--epochs	200	训练轮数
--batch_size	32	批次大小
--lr	0.01	初始学习率（命令行参数，实际代码中使用 Adam lr=0.001）
--n_points	1024	每个点云采样的点数
--seed	42	随机种子
--data_dir	./data	数据集目录

输出

训练完成后会在当前目录生成：

• pct_model.pkl — 训练好的模型权重
• result.json — 测试集预测结果，格式为 {样本编号: 预测类别}

模型结构

输入点云 (B, 3, N)
    │
    ├─ Conv1d(3→128) + BN + ReLU
    ├─ Conv1d(128→128) + BN + ReLU
    │
    ├─ SA_Layer × 4（Self-Attention, 128 channels）
    │
    ├─ Concat [F1, F2, F3, F4] → (B, 512, N)
    ├─ Conv1d(512→1024) + BN + LeakyReLU
    ├─ Global Max Pooling → (B, 1024)
    │
    ├─ FC(1024→512) + BN + ReLU + Dropout(0.5)
    ├─ FC(512→256) + BN + ReLU + Dropout(0.5)
    └─ FC(256→40) → Logits

• 总参数量：约 2.87M
• 优化器：Adam (lr=0.001, weight_decay=1e-4)
• 学习率调度：Cosine Annealing
• 损失函数：Cross Entropy Loss
• 数据增强：随机绕 Y 轴旋转、随机缩放 (±20%)、高斯抖动、随机平移、点序打乱

致谢

• 计图 (Jittor) — 清华大学自研深度学习框架
• ModelNet40 — Princeton ModelNet 数据集
• PCT: Point Cloud Transformer — 原始论文