ModelNet40 PCT 点云分类

本项目基于 Jittor 框架实现 PCT（Point Cloud Transformer）点云分类模型，用于 ModelNet40 三维形状分类任务。脚本会读取预处理后的点云数据，训练分类模型，并 对测试集生成提交文件 result.json。

数据说明

数据文件应放在 data/ 目录下：

• data/train_points.npy：训练集点云
• data/train_labels.npy：训练集标签
• data/test_points.npy：测试集点云
• data/categories.txt：类别名称

每个样本包含 2048 个归一化后的三维点，模型需要预测 40 个类别之一。

文件结构

• pct.py：数据加载、PCT 模型、训练流程和测试集预测导出。
• data/：ModelNet40 点云数据。
• result.json：生成的测试集预测结果，格式为 {样本编号: 预测类别}。
• pct_model.pkl：训练后保存的模型参数。

运行环境

当前机器上的 Jittor 环境位于 conda 环境 cg。

查看可用 conda 环境：

conda env list

快速测试

可以先用较少点数和 1 个 epoch 检查代码是否能完整跑通：

conda run --no-capture-output -n cg python -u pct.py --epochs 1 --batch_size 8 --n_points 128

运行结束后会生成 pct_model.pkl 和 result.json。

完整训练

推荐的 baseline 完整训练命令：

conda run --no-capture-output -n cg python -u pct.py --epochs 200 --batch_size 32 --n_points 1024 --lr 0.01

如果显存允许，也可以使用全部 2048 个点训练，速度会更慢：

conda run --no-capture-output -n cg python -u pct.py --epochs 200 --batch_size 16 --n_points 2048 --lr 0.01

其中 python -u 和 --no-capture-output 用于实时显示日志。第一次运行时，Jittor 可能需要编译 CUDA 算子，开始阶段可能会有几分钟没有训练日志输出。

命令行参数

python pct.py \
  --data_dir ./data \
  --n_points 1024 \
  --batch_size 32 \
  --epochs 200 \
  --lr 0.01 \
  --seed 42

参数含义：

• --data_dir：数据目录，默认 ./data。
• --n_points：每个样本采样的点数。
• --batch_size：训练和预测时的 batch size。
• --epochs：训练轮数。
• --lr：初始学习率。
• --seed：NumPy 和 Jittor 的随机种子。

输出格式

result.json 是一个 JSON 字典，key 为测试样本编号字符串，value 为预测类别整数：

{
  "0": 12,
  "1": 4,
  "2": 36
}

脚本会为测试集全部 2468 个样本生成预测结果。

运行日志与中断

• 训练过程中每 20 个 batch 输出一次 loss、训练准确率和耗时。
• 每个 epoch 结束后输出该 epoch 的平均 loss、训练准确率、学习率和耗时。
• 手动按 Ctrl+C 会立即停止程序。当前脚本只会在训练循环完整结束后保存 pct_model.pkl 并生成 result.json。
• 提供的数据中没有测试集标签，因此无法在本地计算测试集准确率，只能检查 result.json 的格式和样本数量。