PCT_jittor

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

本项目提供了 Point Cloud Transformer (PCT) 的 Jittor 实现，用于 ModelNet40 三维形状分类任务。

主要特性

• 基于 Jittor 框架：利用即时编译（JIT）技术提供高性能的深度学习计算。
• 坐标偏移嵌入：利用线性层（MLP）将原始 3D 坐标映射到高维特征空间。
• **自注意力机制 (Self-Attention)**：堆叠 4 层 SA 层以捕获全局上下文。
• 数据增强：支持随机旋转、随机缩放、高斯抖动及离群点裁剪。

模型架构

1. 输入嵌入层：使用两层 1D 卷积将原始坐标投影至 128 维特征空间。
2. Transformer 编码器：4 层堆叠的 SA 层用于特征精炼，通过偏移注意力（Offset Attention）机制增强鲁棒性。
3. 特征融合：拼接多尺度特征并投影至 1024 维高维空间。
4. 分类头：带有 Dropout 和 BatchNorm 的三层 MLP。

运行环境

• Linux (Ubuntu 推荐)
• Python 3.7+
• Jittor (选定框架)
• NumPy, Scipy, Tqdm

安装核心依赖：

# 安装 Jittor 及系统编译依赖
sudo apt install -y build-essential libomp-dev
pip install jittor numpy scipy tqdm

数据集准备

由于数据集体积较大，未包含在仓库中。请根据以下说明准备数据：

1. 下载 ModelNet40 数据集并预处理为 .npy 格式。
2. 确保文件存放于 data/ 目录：
   • 训练集：data/train_points.npy, data/train_labels.npy
   • 测试集：data/test_points.npy

使用方法

训练模型

使用 GPU 加速运行：

export jt_flags="use_cuda=1"
python3 pct.py --data_dir ./data --batch_size 32 --epochs 200 --lr 0.01

推理预测

训练完成后，脚本会自动对测试集进行推理，并生成 result.json。模型权重将保存为 pct_model.pkl。

贡献与反馈

欢迎通过 Issues 提交反馈或 Pull Requests 改进代码。

代码超参数

参数	默认值	说明
n_points	1024	每个点云采样的点数
batch_size	32	训练批大小
epochs	200	总训练轮数
lr	0.01	初始学习率 (SGD)