PCT_jittor

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

项目简介

本项目基于 Jittor 实现了 Point Cloud Transformer（PCT）点云分类模型，用于 ModelNet40 三维形状分类任务。项目提供了训练、验证式推理和测试集结果导出流程，可直接生成符合提交要求的 result.json。

项目特性

• 基于 Jittor 的 PCT 分类网络实现
• 支持 ModelNet40 点云数据加载与采样
• 内置数据增强策略，包括旋转、缩放、平移和抖动
• 使用余弦退火学习率调度
• 训练完成后自动导出模型权重和测试集预测结果

环境依赖

• Python 3.8+
• Jittor
• NumPy

安装依赖：

pip install jittor numpy

如果已经安装了 Jittor，也可以只补装 NumPy：

pip install numpy

数据集说明

数据文件默认放在 data/ 目录下，脚本会按以下方式读取：

• 训练集：data/train_points.npy 和 data/train_labels.npy
• 测试集：data/test_points.npy
• 类别说明：data/categories.txt

点云输入默认会从每个样本中采样 1024 个点。若数据文件路径不同，可以通过 --data_dir 参数指定。

数据下载与解压

数据集压缩包已放在清华云盘：

https://cloud.tsinghua.edu.cn/d/c1f6be8b2d164c2e940e/

使用方法如下：

1. 下载得到 data.zip。
2. 将 data.zip 放在项目根目录，也就是 pct.py 所在目录。
3. 在项目根目录执行解压命令：

unzip data.zip

解压后会直接生成原来的 data/ 目录，目录结构应与脚本读取方式一致。

项目结构

.
├── pct.py
├── README.md
├── .gitignore
├── pct_model.pkl
├── pct_model_best.pkl
├── result.json

运行方式

直接启动训练与推理：

python pct.py

• --batch_size：批大小，默认 32
• --epochs：训练轮数，默认 200
• --lr：初始学习率，默认 0.001
• --seed：随机种子，默认 42

例如：

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

输出文件

脚本运行结束后会生成以下文件：

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