Jittor-PCT (Jittor)

基于 Jittor 框架的 Point Cloud Transformer (PCT) 实现，用于 ModelNet40 三维形状分类任务。

功能特性

PCT 模型：4 层自注意力（Self-Attention）层 + 特征融合网络
完整数据处理：
- 点云采样（1024 或自定义数量）
- 多种数据增强：随机旋转、缩放、平移、高斯抖动
优化器与学习率调度：
- AdamW 优化器（权重衰减：1e-4）
- 余弦退火（Cosine Annealing）学习率调度
训练与推理：批处理训练、逐 epoch 性能统计、测试集预测导出

项目结构

.
├─ pct.py                 # 完整的训练 + 推理脚本
├─ data/
│  ├─ train_points.npy    # 训练集点云 (N, 2048, 3)
│  ├─ train_labels.npy    # 训练集标签 (N,)
│  ├─ test_points.npy     # 测试集点云 (N, 2048, 3)
│  └─ categories.txt      # 40 个类别名（可选）
├─ dockerfile             # CUDA 12.2 + Python 3.8 Docker 镜像
├─ start.sh               # Docker 容器启动脚本
├─ pct_model.pkl          # 训练后的模型参数（自动生成）
├─ result.json            # 测试集预测结果（自动生成）
└─ README.md              # 本文档

环境要求

Python：3.8+
依赖：Jittor、NumPy
硬件：CUDA GPU（可选，默认使用 GPU；不用 GPU 见下文）

安装依赖

pip install jittor numpy

如不使用 GPU

编辑 pct.py 第 283 行：

jt.flags.use_cuda = 0  # 改为 0 使用 CPU

使用方法

基础用法

python pct.py

使用默认参数训练 250 epochs，输出 result.json。

自定义参数

python pct.py \
    --data_dir ./data \           # 数据集路径
    --n_points 1024 \             # 每样本点数采样
    --batch_size 64 \             # 批大小
    --epochs 250 \                # 训练轮数
    --lr 0.001 \                  # 初始学习率
    --seed 42                     # 随机种子

参数说明

参数	默认值	说明
`--data_dir`	`./data`	ModelNet40 数据集存放目录
`--n_points`	`1024`	从每个点云采样的点数（原始 2048）
`--batch_size`	`64`	训练批大小
`--epochs`	`250`	总训练轮数
`--lr`	`0.001`	AdamW 初始学习率
`--seed`	`42`	全局随机种子

输出文件

训练完成后生成：

**pct_model.pkl**：模型参数文件，可用 model.load() 加载
**result.json**：格式为 {"样本id": 预测类别, ...}，共 40 个类别（0-39）

示例：

{
  "0": 5,
  "1": 12,
  "2": 8,
  ...
}

数据增强

在 ModelNet40Dataset.__getitem__() 中实现，训练时自动应用，包括：

随机旋转：绕 Y 轴旋转 [0, 2π]
随机缩放：缩放因子 [0.8, 1.25]
随机平移：平移范围 [-0.1, 0.1]
高斯抖动：标准差 0.01，裁剪 [-0.05, 0.05]

测试集不使用数据增强。

模型架构

输入: (B, 3, 1024) 点云坐标
↓
Conv1d (3 → 128) + ReLU + BatchNorm
Conv1d (128 → 128) + ReLU + BatchNorm
↓
4 × Self-Attention Layer (128 channels)
↓
Concat (128×4 → 512) + Conv1d (512 → 1024) + LeakyReLU
↓
Global Max Pooling → (B, 1024)
↓
FC (1024 → 512) + ReLU + Dropout
FC (512 → 256) + ReLU + Dropout
FC (256 → 40) [Logits]
↓
输出: (B, 40) 类别 logits

模型参数量约为几百万，具体数值训练时输出。

Docker 支持（可选）

构建镜像

docker build -t my-jittor-cuda:12.2 -f dockerfile .

运行容器

bash start.sh

或手动运行：

docker run -it --rm --gpus all \
    -v /path/to/PA3:/workspace \
    -v /path/to/jittor_cache:/root/.cache/jittor \
    my-jittor-cuda:12.2 bash