PCT_DGCNN_jittor

基于 Jittor 实现的 ModelNet40 点云分类项目。项目从 PA3 提供的 PCT baseline 出发，补齐训练、预测和提交流程，并进一步实现了 DGCNN / EdgeConv 模型。最终采用 DGCNN best176 版本，线上准确率达到 **89.3%**。

项目简介

本项目用于完成 ModelNet40 三维点云分类任务。输入为点云坐标，输出为 40 个类别中的一个类别编号。项目保留了作业原始 PCT 模型，同时增加了 DGCNN 模型，重点增强局部几何结构建模能力。

开发这个项目的主要目标有三个：

完成 Jittor 框架下的点云分类训练与预测流程；
生成符合评测平台要求的 result.json；
在 baseline 基础上通过 DGCNN、label smoothing 和 test-time voting 提升泛化性能。

主要特性

支持 pct 和 dgcnn 两种模型；
使用 DGCNN / EdgeConv 建模点云局部邻域关系；
支持随机绕 Y 轴旋转的数据增强；
支持 label smoothing、dropout 和 cosine learning rate schedule；
支持 checkpoint 保存和断点续训；
支持 test-time voting，降低随机采样带来的预测波动；
自动导出并校验平台需要的 result.json。

目录结构

PCT_DGCNN_jittor/
  README.md
  .gitignore
  requirements.txt
  pct.py
  pct_model.pkl
  result.json

说明：

pct.py：主代码，包含数据集、PCT、DGCNN、训练、预测和导出逻辑；
pct_model.pkl：最终模型权重；
result.json：最终测试集预测结果；
requirements.txt：Python 依赖；
data/ 不放入仓库，需要自行准备。

环境配置

推荐使用 WSL Ubuntu + conda：

conda create -n jittor-pct python=3.10 -y
conda activate jittor-pct
python -m pip install -U pip
python -m pip install -r requirements.txt
python -m jittor.test.test_example
python -m jittor.test.test_cuda

如果 CUDA 测试失败，可以先使用 CPU 做 smoke test，或者按照 Jittor 官方文档配置 CUDA 环境。

数据准备

请将数据放到项目根目录的 data/ 文件夹下：

data/
  train_points.npy
  train_labels.npy
  test_points.npy

数据规模：

train_points.npy: (9843, 2048, 3)
test_points.npy:  (2468, 2048, 3)
classes:          0..39

快速使用

如果已经有训练好的 pct_model.pkl，可以直接运行预测并生成 result.json：

python -u pct.py \
  --model dgcnn \
  --epochs 0 \
  --batch_size 8 \
  --n_points 1024 \
  --use_cuda 1 \
  --model_path pct_model.pkl \
  --result_path result.json \
  --vote_num 10

训练命令

DGCNN 正式训练命令示例：

python -u pct.py \
  --model dgcnn \
  --epochs 200 \
  --batch_size 8 \
  --n_points 1024 \
  --lr 0.001 \
  --optimizer adam \
  --label_smoothing 0.1 \
  --k 20 \
  --dropout 0.5 \
  --use_cuda 1 \
  --num_workers 2 \
  --predict_num_workers 0 \
  --checkpoint_interval 5 \
  --checkpoint_dir checkpoints_dgcnn_adam \
  --model_path pct_model.pkl \
  --best_model_path best_pct_model.pkl \
  --result_path result.json \
  --log_interval 50 \
  --vote_num 10

如果显存不足，可以将 --batch_size 改为 4。

断点续训

训练过程中会定期保存 checkpoint。中断后可以使用 --resume 继续训练：

python -u pct.py \
  --model dgcnn \
  --epochs 200 \
  --batch_size 8 \
  --n_points 1024 \
  --lr 0.001 \
  --optimizer adam \
  --label_smoothing 0.1 \
  --k 20 \
  --dropout 0.5 \
  --use_cuda 1 \
  --resume checkpoints_dgcnn_adam/checkpoint_epoch_160.pkl \
  --checkpoint_dir checkpoints_dgcnn_adam_resume \
  --model_path pct_model.pkl \
  --best_model_path best_pct_model.pkl \
  --result_path result.json \
  --vote_num 10

结果提交

评测平台要求 result.json 的格式为：

{
  "0": 12,
  "1": 3,
  "2": 37
}

其中 key 是测试样本编号字符串，value 是 0..39 的整数类别。

校验命令：

python -c "import json; r=json.load(open('result.json')); assert len(r)==2468; assert set(r)==set(map(str, range(2468))); assert all(isinstance(v,int) and 0<=v<40 for v in r.values()); print('OK')"

生成提交压缩包时，压缩包内部文件名必须是 result.json。

实验结果

最终选择 DGCNN best176 版本：

model: DGCNN
n_points: 1024
k: 20
optimizer: Adam
label_smoothing: 0.1
dropout: 0.5
vote_num: 10
online accuracy: 89.3%

相比原始 PCT baseline，DGCNN 的主要优势是动态构建邻域图，并通过 EdgeConv 学习局部几何关系。这对 ModelNet40 中的椅子、桌子、飞机等三维物体分类更有效。

参考

Jittor: https://github.com/Jittor/jittor
DGCNN: Dynamic Graph CNN for Learning on Point Clouds
ModelNet40 point cloud classification task

License

本项目主要用于课程作业与学习交流。如需进一步开源发布，可根据课程要求补充正式开源许可证。