PCT-Jittor ModelNet40 Classification

基于 Jittor 实现的 PCT（Point Cloud Transformer）点云分类模型，用于 ModelNet40 三维形状分类任务。代码包含数据加载、点云数据增强、PCT 模型定义、训练流程、余弦学习率调度以及测试集预测结果导出。

项目结构

.
├── pct.py          # 训练、预测和模型定义
├── README.md      # 项目说明
└── .gitignore     # 忽略数据集、模型权重和运行产物

数据集.zip压缩包请在https://cloud.tsinghua.edu.cn/f/f003de5a2e914d1e9e0e/?dl=1 下载 在训练时，请将训练数据集解压到 data/ 目录：

data/
├── train_points.npy
├── train_labels.npy
├── test_points.npy
└── categories.txt

环境要求

建议使用 Python 3.8+，并安装：

pip install numpy jittor

如需使用 GPU，请先确认本机 CUDA 环境可被 Jittor 正常识别。代码中默认启用：

jt.flags.use_cuda = 1

快速开始

在仓库根目录运行：

python pct.py --data_dir ./data

默认训练配置：

n_points   = 2048
batch_size = 16
epochs     = 200
lr         = 0.01
eval_votes = 10

训练完成后会生成：

pct_model.pkl   # 训练好的模型权重
result.json     # 测试集预测结果

result.json 的格式如下：

{
  "0": 4,
  "1": 35,
  "2": 10
}

其中 key 为测试集样本编号，value 为预测类别编号，范围为 0-39。

常用命令

完整精度训练：

python pct.py --data_dir ./data --n_points 2048 --batch_size 16 --epochs 200 --lr 0.01 --eval_votes 10

显存较小时可以减小 batch size：

python pct.py --data_dir ./data --n_points 2048 --batch_size 8 --epochs 200 --lr 0.01

快速调试可以减少点数和训练轮数：

python pct.py --data_dir ./data --n_points 1024 --batch_size 32 --epochs 20 --lr 0.01 --eval_votes 1

方法简介

本项目使用 PCT 结构处理点云分类任务：

• 输入为每个样本的三维点云，形状为 (N, 3)。
• 通过 Conv1d + BatchNorm + ReLU 提取局部点特征。
• 使用多层 self-attention 建模点与点之间的关系。
• 通过全局池化得到形状级特征。
• 使用全连接分类头输出 40 个类别的预测 logits。

训练阶段使用的增强策略包括：

• 随机绕 Y 轴旋转
• 随机缩放
• 轻微 jitter 抖动
• 随机打乱点顺序

优化器使用 SGD，带 momentum 和 weight decay，并配合 cosine annealing 学习率调度。

提交与开源说明

仓库中的 .gitignore 已忽略数据集、模型权重和运行输出，包括：

data/
data.zip
*.npy
*.pkl
result.json
result.zip
__pycache__/