PCT_jittor

本仓库为第六届计图（Jittor）人工智能挑战赛赛道二（点云分类）热身赛的参赛实现，基于 Jittor 框架。

最终在赛道二评测集上的分类准确率为 0.8489。

模型结构

PCT 分类网络由以下模块串联而成：

Input Embedding：两层 Conv1d 将每个点的三维坐标映射到 128 维特征（ $3 \to 128 \to 128$ ）。
Offset Self-Attention × 4：四个共享 $W_{Q} = W_{K}$ 权重的自注意力模块串联，每层均输出 128 维特征。
Concat + Conv-Fuse：拼接四层自注意力的输出得到 512 维特征，再经一层 Conv1d（LeakyReLU）提升至 1024 维。
Max-Pool：沿点维度做最大池化，得到 1024 维全局描述子。
MLP Classifier：三层全连接 $1024 \to 512 \to 256 \to 40$ ，中间层配 Dropout（丢弃率 0.5）。

模型参数量约 1.37 M。

训练设置

超参数	取值
输入点数	1024（模型 A）/ 2048（模型 B）
Batch size	32
Epoch	400（模型 A）/ 300（模型 B）
Optimizer	AdamW，weight decay $= 10^{-4}$
Learning rate	$10^{-3} \to 10^{-5}$ ，余弦退火
Loss	Cross-entropy + label smoothing ( $\epsilon = 0.1$ )
Gradient clipping	Global $L_{2}$ norm $\le 1.0$
数据增强	绕 $Y$ 轴随机旋转、各向同性缩放 $[0.8, 1.25]$ 、平移 $\pm 0.1$ 、高斯抖动 ( $\sigma = 0.01$ , clip $\pm 0.05$ )
测试阶段	每模型 15 次 test-time augmentation (TTA) + 双模型 softmax 概率平均

项目结构

.
├── pct.py            # 数据集、模型、训练、单模型推理
├── ensemble.py       # 双模型 logits 集成 + 测试时增强
├── eval_train.py     # 在干净训练集上评估模型精度
├── pyproject.toml    # 依赖声明（uv 管理）
├── uv.lock           # 依赖锁定
├── README.md
└── .gitignore

安装

依赖：Python 3.10、g++ 12、CUDA 显卡（本仓库在 NVIDIA RTX 4090 + CUDA 12.2 环境下测试通过）。推荐使用 uv 管理虚拟环境：

uv sync
uv run python -m jittor_utils.install_cuda

数据集

ModelNet40 已预处理为四个文件，可从挑战赛页面下载。

运行

训练单模型

以模型 A（输入点数 1024）的默认配置训练并产出 result.json：

uv run python pct.py

训练双模型并集成

# 1) 模型 A：输入点数 1024
uv run python pct.py \
    --n_points 1024 --batch_size 32 --epochs 400 \
    --label_smoothing 0.1 --model_path pct_model_n1024.pkl

# 2) 模型 B：输入点数 2048
uv run python pct.py \
    --n_points 2048 --batch_size 32 --epochs 300 \
    --label_smoothing 0.1 --model_path pct_model_n2048.pkl

# 3) 集成两模型生成最终 result.json
uv run python ensemble.py \
    --model_a pct_model_n1024.pkl --npts_a 1024 \
    --model_b pct_model_n2048.pkl --npts_b 2048 \
    --n_votes 15

常用参数

参数	默认值	说明
`--data_dir`	`../data`	数据集所在目录
`--n_points`	1024	每个样本采样的点数
`--batch_size`	32	训练与推理的 batch size
`--epochs`	400	训练 epoch 数
`--lr`	$10^{-3}$	初始 learning rate
`--optimizer`	`adamw`	Optimizer（可选 `sgd` / `adamw`）
`--label_smoothing`	0.1	Label smoothing 系数
`--grad_clip`	1.0	Gradient clipping 的 $L_{2}$ norm 上限
`--n_votes`	10	TTA 投票次数
`--model_path`	`pct_model.pkl`	模型权重保存路径
`--result_path`	`result.json`	测试集预测结果保存路径

输出

pct_model.pkl — 训练完成的 PCT 模型权重（Jittor pickle 格式）。
result.json — 测试集每个样本的预测类别，格式为 {"样本编号": 类别 id} 的 JSON 字典。

致谢

网络结构参考 Guo et al., PCT: Point Cloud Transformer, Computational Visual Media 2021。
训练框架基于 Jittor。