PCT_jittor

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

本项目是图形学实验 PA3 的点云分类实现，基于 Jittor 深度学习框架，在 ModelNet40 数据集上训练 Point Cloud Transformer，并生成测试集预测文件 result.json。

Features

• 使用 Jittor 实现 PCT 点云分类网络；
• 支持 ModelNet40 预处理点云数据读取；
• 支持随机采样、随机旋转、缩放、平移、抖动等点云数据增强；
• 支持 Adam / SGD 优化器和 Cosine Annealing 学习率调度；
• 支持 checkpoint 保存、模型恢复训练和测试集预测；
• 支持 test-time voting，提高预测稳定性。

Project Structure

.
├── pct.py          # 主程序：数据读取、模型定义、训练和预测
├── README.md      # 项目说明
└── .gitignore     # Git 忽略规则

数据集、模型权重和预测结果文件体积较大或由程序生成，默认不纳入 Git 仓库。运行时请将数据放在 data/ 目录下。

Requirements

• Python 3.8+
• Jittor
• NumPy
• C++ compiler, such as g++
• CUDA GPU is recommended for training PCT

安装依赖：

pip install jittor numpy

检查 GPU：

nvidia-smi

Dataset

本实验使用预处理后的 ModelNet40 点云数据。目录结构如下：

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

其中：

• train_points.npy：训练集点云，形状为 (N, 2048, 3)；
• train_labels.npy：训练集类别标签；
• test_points.npy：测试集点云，无标签；
• categories.txt：类别名称列表。

数据集下载链接：https://cloud.tsinghua.edu.cn/d/53154af941ba49ba99f1/

Usage

Train PCT

GPU 环境下推荐使用 PCT 模型训练：

python pct.py \
  --model pct \
  --data_dir ./data \
  --optimizer adam \
  --lr 0.001 \
  --epochs 200 \
  --n_points 1024 \
  --batch_size 16 \
  --test_votes 5 \
  --use_cuda

训练完成后，程序会保存模型文件并生成预测文件：

pct_model.pkl
result.json

Fine-tune

可以从已有模型继续微调：

python pct.py \
  --model pct \
  --data_dir ./data \
  --optimizer adam \
  --lr 0.0002 \
  --epochs 80 \
  --n_points 1024 \
  --batch_size 16 \
  --test_votes 7 \
  --use_cuda \
  --resume pct_model.pkl \
  --save_path pct_adam_ft.pkl \
  --result_path result.json

Evaluation-only Prediction

如果已经有训练好的模型，可以只生成测试集预测结果：

python pct.py \
  --model pct \
  --data_dir ./data \
  --resume pct_model.pkl \
  --eval_only \
  --n_points 1024 \
  --batch_size 16 \
  --test_votes 7 \
  --use_cuda \
  --result_path result.json

CPU-friendly Option

PCT 的自注意力模块需要构造 N x N 注意力矩阵，在 CPU 上训练较慢。本项目也保留了一个轻量 PointNet 分类器，便于调试代码流程：

python pct.py \
  --model pointnet \
  --data_dir ./data \
  --optimizer adam \
  --lr 0.001 \
  --epochs 120 \
  --n_points 1024 \
  --batch_size 16

注意：最终提交使用的是 PCT 模型结果。

Main Arguments

Argument	Description
--model	模型类型，可选 pct 或 pointnet
--data_dir	数据集目录
--n_points	每个点云采样点数
--batch_size	批大小
--epochs	训练轮数
--optimizer	优化器，可选 adam 或 sgd
--lr	初始学习率
--test_votes	测试阶段投票次数
--use_cuda	使用 CUDA GPU
--resume	从已有模型继续训练或预测
--eval_only	只进行预测，不训练
--save_path	模型保存路径
--result_path	预测结果保存路径

Result Format

输出文件 result.json 为字典格式，键为测试样本编号，值为预测类别编号：

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

该文件可直接提交到课程评测平台进行 Accuracy 评测。