PCT-Jittor-ModelNet40

基于 Jittor 实现的 Point Cloud Transformer（PCT）点云分类项目，用于 ModelNet40 三维形状分类任务。本项目包含点云数据读取、训练数据增强、PCT 分类网络、训练流程、模型保存和测试集预测结果导出。

项目简介

本项目是图形学实验 PA3 的代码实现，目标是在 ModelNet40 点云数据集上训练一个 PCT 分类模型。模型输入为三维点云坐标，输出 40 个类别的分类 logits，并在测试集上生成 result.json 预测文件。

项目使用的主要框架和工具：

Jittor：深度学习框架
Python：训练与推理脚本
NumPy：点云数据读取与数据增强
ModelNet40：三维形状分类数据集

项目特点

使用 Jittor 搭建 Point Cloud Transformer（PCT）分类网络。
支持 ModelNet40 预处理点云数据读取。
训练阶段包含随机绕 Y 轴旋转、随机缩放和随机抖动等数据增强。
支持 SGD 或 Adam 优化器，默认使用 SGD 配合余弦退火学习率调度。
支持快速 smoke 测试，便于本地检查完整训练和预测流程。
训练完成后保存模型文件，并导出测试集预测结果。

目录结构

.
├── pct.py              # PCT 模型、数据集、训练和预测主程序
├── README.md           # 项目说明文件
├── .gitignore          # Git 忽略规则
├── docs/               # 作业说明与报告相关文档
└── data/               # 数据集目录，本地使用，不建议提交到开源仓库

数据目录在本地应包含：

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

环境要求

建议在 WSL 中使用 conda 环境运行。本项目本地测试环境为：

WSL
conda 环境：jittor
Python
Jittor

进入环境：

conda activate jittor

如果从 Windows PowerShell 调用 WSL，可使用：

wsl bash -lc 'source /root/miniconda3/etc/profile.d/conda.sh && conda activate jittor && cd "/mnt/d/Users/25931/University/大二下/计图/作业/PA3" && python pct.py --help'

快速测试

使用 smoke 模式可以快速检查数据读取、前向传播、训练、模型保存和预测导出流程：

python pct.py --smoke --no_cuda --model_path smoke_pct_model.pkl --output smoke_result.json

该命令会使用少量样本进行 1 个 epoch 的快速测试，并生成：

smoke_pct_model.pkl
smoke_result.json

正式训练与预测

默认训练命令：

python pct.py

默认配置：

点数：1024
batch size：32
epoch：200
学习率：0.01
优化器：sgd
模型输出：pct_model.pkl
预测结果：result.json

训练完成后，脚本会自动保存模型并在测试集上生成预测文件。

常用参数

python pct.py \
  --data_dir ./data \
  --n_points 1024 \
  --batch_size 32 \
  --epochs 200 \
  --lr 0.01 \
  --optimizer sgd \
  --model_path pct_model.pkl \
  --output result.json

参数说明：

--data_dir：数据集目录，默认 ./data
--n_points：每个点云样本采样点数，默认 1024
--batch_size：批大小，默认 32
--epochs：训练轮数，默认 200
--lr：学习率，默认 0.01
--optimizer：优化器，可选 sgd 或 adam，默认 sgd
--model_path：模型保存路径，默认 pct_model.pkl
--output：预测结果保存路径，默认 result.json
--no_cuda：使用 CPU 运行，适合快速调试
--smoke：快速流程测试模式

模型结构

当前 PCT 模型输入为 (B, 3, N) 的点云数据，整体结构如下：

输入点云
  -> Conv1d(3, 128) + BatchNorm + ReLU
  -> Conv1d(128, 128) + BatchNorm + ReLU
  -> SA_Layer × 4
  -> 拼接四个注意力层输出，得到 512 通道特征
  -> Conv1d(512, 1024) + BatchNorm + LeakyReLU
  -> 点维度最大池化，得到 1024 维全局特征
  -> Linear(1024, 512) + BatchNorm + ReLU + Dropout
  -> Linear(512, 256) + BatchNorm + ReLU + Dropout
  -> Linear(256, 40)

其中 SA_Layer 使用 Query、Key、Value 计算点与点之间的自注意力关系，并通过残差连接增强全局几何特征建模能力。

输出文件

正式训练完成后会生成：

pct_model.pkl：训练得到的 Jittor 模型文件
result.json：测试集预测结果，格式为 {样本编号: 预测类别}

这些文件属于训练产物，不建议提交到 Gitlink 开源仓库。提交课程作业时可按课程要求单独打包。

Gitlink 开源说明

建议在 Gitlink 中使用以下项目信息：

项目名称：PCT-Jittor-ModelNet40
项目简介：基于 Jittor 实现的 Point Cloud Transformer（PCT）点云分类项目，用于 ModelNet40 三维形状分类任务。

开源仓库建议包含：

pct.py
README.md
.gitignore
docs/ 中的作业说明或报告文档

不建议提交：

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

参考

Point Cloud Transformer（PCT）
Jittor 深度学习框架
ModelNet40 三维形状分类数据集