PCT_jittor

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

本项目是论文 PCT: Point Cloud Transformer 在 Jittor 框架下的复现，用于完成 ModelNet40 三维形状分类任务。

项目特性

✅ Jittor 原生实现，使用 jt.nn 与 Dataset 接口，自动 GPU 加速。
✅ PCT 论文版 Offset-Attention：Q/K 独立参数 + 双重归一化 + 残差。
✅ 完整数据增强：随机绕 Y 轴旋转、各向异性缩放、平移、高斯抖动、点序 shuffle。
✅ 点云归一化：中心化 + 缩放到单位球。
✅ 训练集自动 9:1 切分验证集，用 val_acc 选 best 模型。
✅ Label Smoothing 交叉熵 + SGD + Warmup + Cosine Annealing 学习率调度。
✅ 测试时 TTA（Test-Time Augmentation）：多次随机采样取平均 logits。

环境依赖

依赖	版本建议
Python	≥ 3.7
Jittor	≥ 1.3.0
NumPy	≥ 1.19
CUDA	推荐 10.2 / 11.x（CPU 也可运行）

安装：

pip install jittor numpy
# 可选：验证 Jittor 安装
python -m jittor.test.test_example

数据集

本项目使用 ModelNet40 点云数据集（共 40 类 CAD 模型），数据格式为 numpy 数组：

文件	形状	说明
`train_points.npy`	`(N_train, 2048, 3)`	训练集点云
`train_labels.npy`	`(N_train,)`	训练集类别标签（0~39）
`test_points.npy`	`(N_test, 2048, 3)`	测试集点云（无标签）
`categories.txt`	文本	40 个类别名（按标签序号排列）

⚠️ 由于数据集体积较大，未上传至本仓库。请通过以下方式获取： 下载地址：https://www.educoder.net/competitions/Jittor-7 放置后目录应为：

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

项目结构

PCT_jittor/
├── pct.py                  # 主程序（数据集 / 模型 / 训练 / 推理 一体）
├── data/
│   └── data/
│       └── categories.txt  # 40 个类别名（标签索引对照）
├── README.md
└── .gitignore

快速开始

1. 训练 + 自动推理（推荐）

python pct.py \
    --data_dir ./data/data \
    --epochs 250 \
    --batch_size 32 \
    --lr 0.01 \
    --optimizer sgd \
    --tta 3

训练完成后会自动：

保存最佳模型到 pct_model.pkl、最后一轮模型到 pct_model_last.pkl；
加载最佳模型对测试集做 TTA 推理；
输出预测结果到 result.json（格式：{"样本编号": 类别索引, ...}）。

2. 仅推理（已有训练权重）

将训练好的 pct_model.pkl 放到项目根目录，然后：

python pct.py --epochs 0 --tta 3

提示：当前主程序未单独提供「跳过训练」开关。若只想做推理，可以将 --epochs 设为很小的值，或自行在 main() 中跳过训练循环；推理逻辑见 predict() 函数。

训练参数说明

参数	默认值	说明
`--data_dir`	`./data/data`	数据目录
`--n_points`	`1024`	每个样本采样点数
`--batch_size`	`32`	批大小
`--epochs`	`250`	训练轮数
`--lr`	`0.01`	初始学习率
`--weight_decay`	`1e-4`	权重衰减
`--warmup`	`10`	学习率 warmup 轮数
`--smoothing`	`0.1`	Label smoothing 系数
`--optimizer`	`sgd`	优化器：`sgd` / `adam` / `adamw`
`--val_ratio`	`0.1`	验证集划分比例（0 表示用全部训练集）
`--tta`	`3`	测试时增强次数
`--out_model`	`pct_model.pkl`	最佳模型保存路径
`--out_json`	`result.json`	预测结果保存路径
`--seed`	`42`	随机种子
`--num_workers`	`4`	DataLoader 工作进程数
`--no_cuda`	—	添加该 flag 则使用 CPU

实验结果

模型	数据集	框架	准确率 (Acc)
PCT (本项目)	ModelNet40	Jittor	≈ ~86%（具体取决于训练轮数 / 随机种子）

模型参数量：约 2.88 M。

预测结果文件 result.json 示例：

{
  "0": 0,
  "1": 0,
  "2": 0
}

其中 key 为测试样本索引（与 test_points.npy 中顺序一致），value 为类别索引（对应 categories.txt 中的行号）。

模型架构

Input (B, 3, N)
   │
   ├─► Conv1d(3→128) + BN + ReLU
   └─► Conv1d(128→128) + BN + ReLU                # Embedding
            │
   ┌────────┴────────┐
   │ SA_Layer (×4)   │  # Offset-Attention，输出 4 份 (B,128,N)
   └────────┬────────┘
            ▼
       Concat → (B, 512, N)
            │
   Conv1d(512→1024) + BN + LeakyReLU              # Fuse
            │
       MaxPool over N → (B, 1024)
            │
   Linear(1024→512) + BN + ReLU + Dropout
   Linear(512→256)  + BN + ReLU + Dropout
   Linear(256→40)
            │
            ▼
       Logits (B, 40)

SA_Layer 中的 offset-attention：

attn = softmax(Q·K)  →  按列归一化
x_r  = V · attn
out  = x + LBR(trans_conv(x - x_r))   # 残差形式

致谢与参考

原论文：Guo et al., PCT: Point Cloud Transformer, Computational Visual Media 2021.
Jittor 框架