PCT 点云分类 — 第六届计图挑战赛热身赛二

基于 PCT (Point Cloud Transformer) 和 Jittor 框架的 ModelNet40 三维点云分类项目。

赛题简介

在 ModelNet40 数据集上训练点云分类模型，对测试集中的 2468 个点云样本预测其所属类别（共 40 类）。测试集准确率 ≥ 80% 即为通过。

项目结构

.
├── main.py                  # 主入口（训练 / 推理）
├── baseline_code.py         # 官方 baseline 参考代码
├── requirements.txt         # Python 依赖
├── data/                    # 数据目录
│   ├── train_points.npy     # 训练集点云 (9843, 2048, 3)
│   ├── train_labels.npy     # 训练集标签 (9843,)  值域 0-39
│   ├── test_points.npy      # 测试集点云 (2468, 2048, 3)
│   └── categories.txt       # 40 个类别名称
├── models/
│   └── pct.py               # PCT 模型定义（SA_Layer + PCTClassifier）
├── utils/
│   ├── dataset.py           # ModelNet40 数据集加载器
│   ├── augmentor.py         # 数据增强（旋转、缩放、平移、抖动）
│   ├── trainer.py           # 训练管理器（含早停、余弦退火 LR）
│   ├── inferencer.py        # 推理模块
│   └── result_generator.py  # 生成 result.json 并打包 result.zip
├── checkpoints/             # 模型检查点保存目录
├── survey/                  # 相关论文调研笔记
└── tasks/                   # 赛题说明

环境配置

1. 创建 Conda 环境

conda create -n jittor python=3.9 -y
conda activate jittor

2. 安装编译依赖（Linux）

conda install -c conda-forge gcc=10 gxx=10 -y
conda install -c conda-forge libgomp -y

Windows 用户需确保已安装 MSVC 或 MinGW 编译工具链。

3. 安装 Python 依赖

pip install -r requirements.txt

requirements.txt 内容：

jittor
numpy

4. 准备数据

将数据文件放入 data/ 目录下，确保包含以下文件：

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

数据下载地址：https://cloud.tsinghua.edu.cn/f/f003de5a2e914d1e9e0e/?dl=1

5. 确认 GPU 可用

Jittor 默认使用 CUDA，请确保已安装 NVIDIA 驱动和 CUDA Toolkit。首次运行时 Jittor 会进行 JIT 编译，可能需要等待 1-5 分钟。

快速开始

训练

python main.py --mode train --epochs 250 --lr 0.01

常用参数：

参数	默认值	说明
`--mode`	`train`	运行模式：`train` 或 `test`
`--epochs`	`250`	训练轮数
`--lr`	`0.01`	初始学习率
`--batch_size`	`16`	批大小
`--n_points`	`512`	每个样本采样点数（显存不足时减小）
`--patience`	`20`	早停耐心值（连续 N 轮无提升则停止）
`--checkpoint_dir`	`checkpoints`	检查点保存目录
`--seed`	`42`	随机种子

训练过程中会自动保存：

checkpoints/best_model.pkl — 最佳模型（训练准确率最高）
checkpoints/last_model.pkl — 最终模型

推理 & 生成提交文件

python main.py --mode test --checkpoint checkpoints/best_model.pkl

推理完成后会在项目根目录生成：

result.json — 预测结果
result.zip — 可直接提交的压缩包

提交

将生成的 result.zip 上传到比赛平台即可。

模型架构

PCTClassifier 基于 Point Cloud Transformer：

Input Embedding：两层 1D 卷积将 3D 坐标映射到 128 维特征
Self-Attention × 4：4 层 SA_Layer 逐层提取全局上下文特征
Feature Fusion：拼接 4 层输出 (512D) → 1D 卷积融合到 1024D
Global Pooling：Max Pooling 聚合全局特征
Classification Head：1024 → 512 → 256 → 40，含 BN + Dropout

数据增强策略

训练时自动应用以下增强（见 utils/augmentor.py）：

随机旋转：绕 Y 轴随机旋转 [0, 2π]
随机缩放：缩放因子 [0.8, 1.2]
随机平移：每轴独立平移 [-0.1, 0.1]
高斯抖动：σ=0.01，截断到 [-0.05, 0.05]

优化方向

调整超参数（学习率、batch size、训练轮数）
增强数据增强策略（多轴旋转、随机丢点等）
改进模型结构（增加注意力头数、调整隐藏层维度）
使用更强的模型架构（PointNet++、DGCNN 等）

提交格式

result.zip
└── result.json

result.json 示例：

{
  "0": 4,
  "1": 35,
  "2": 10,
  ...
}

key 为样本编号（字符串，从 “0” 开始），value 为预测类别（整数，0-39）。