PCT_jittor

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

本项目使用 Jittor 复现 PCT（Point Cloud Transformer）分类网络，用于完成 ModelNet40 三维点云分类任务。模型读取预处理后的 .npy 点云数据，训练分类器，并在无标签测试集上生成符合提交要求的 result.json。

项目功能

使用 Jittor 实现 PCT 点云分类模型。
支持 ModelNet40 训练集、验证集划分和测试集预测。
支持基于 YAML 的实验配置。
支持点云数据增强，包括绕 Y 轴旋转、小角度扰动、随机缩放等。
支持 SGD / Adam / AdamW 优化器配置。
支持 warmup cosine / cosine 学习率调度。
保存 Jittor .pkl 模型 checkpoint，并生成测试集预测文件 result.json。

项目结构

pct/
├── pct.py                     # Jittor 版 PCT 训练、验证、预测主程序
├── requirements.txt           # Python 依赖
├── README.md                  # 项目说明
├── result.json                # 测试集预测结果
├── result.zip                 # 提交用预测压缩包
├── config/
│   └── config_scratch.yaml    # 当前训练/预测配置
├── data/
│   ├── train_points.npy       # 训练点云，形状通常为 (N, 2048, 3)
│   ├── train_labels.npy       # 训练标签，形状通常为 (N,)
│   ├── test_points.npy        # 测试点云，无标签
│   └── categories.txt         # ModelNet40 类别名称
└── snapshots/
    └── best_model.pkl         # Jittor 保存的最佳模型参数

环境依赖

Jittor 当前主要支持 Linux / Windows / WSL 环境，并需要可用的 C++ 编译器。GPU 训练需要 CUDA 环境。

推荐使用 Conda 创建独立环境：

conda create -n jittor python=3.12 -y
conda activate jittor
pip install -r requirements.txt

如果只想手动安装最小依赖：

pip install jittor numpy pyyaml tensorboardX

数据格式

当前推荐使用的配置文件是 config/config_scratch.yaml，其中 data.data_dir 指向 data/。

训练和预测需要以下文件：

data/train_points.npy
data/train_labels.npy
data/test_points.npy
data/categories.txt

数据集类 ModelNet40Dataset 会从每个点云中采样 n_points 个点。当前配置为：

data:
  n_points: 1024
  num_classes: 40
  val_split: 0.1

训练时会将训练集按随机种子划分为训练集和验证集；测试集没有标签，预测时返回 {样本编号: 预测类别} 格式的 JSON。

快速开始

训练模型：

python pct.py train --config config/config_scratch.yaml

仅生成测试集预测：

python pct.py test --config config/config_scratch.yaml

训练完成后，程序会自动加载最佳模型并生成：

result.json

如果需要提交压缩包，可以将 result.json 压缩为 result.zip。

配置说明

本项目不在 pct.py 中维护默认配置，所有实验参数都写在 config/config_scratch.yaml 中.

主要配置段如下：

run:
  mode: train
  seed: 42
  device: auto

data:
  data_dir: data
  num_classes: 40
  n_points: 1024
  val_split: 0.1

train:
  epochs: 200
  batch_size: 32
  num_workers: 4

run.mode 可以是：

train：训练模型、保存 checkpoint，并生成测试集预测。
test：加载已有模型，只生成测试集预测。

run.device 可以是：

auto：自动使用 CUDA，如果 Jittor 检测不到 CUDA 则使用 CPU。
cuda：强制使用 CUDA。
cpu：强制使用 CPU。

模型与训练策略

模型主体在 pct.py 中实现，包含以下模块：

两层 Conv1d + BatchNorm1d + ReLU 初始特征提取。
多层 SA_Layer 自注意力模块。
拼接所有 self-attention 层输出后，通过 Conv1d 融合到 1024 维特征。
对点维度做 max pooling 得到全局特征。
使用三层全连接分类头输出 40 类 logits。

当前配置中的模型参数：

model:
  channels: 128
  attention_layers: 4
  attention_reduction: 4
  share_qk: true
  fusion_channels: 1024
  fc_dims: [512, 256]
  dropout: [0.5, 0.5]

当前训练策略：

optimizer:
  type: sgd
  lr: 0.01
  momentum: 0.9
  weight_decay: 0.0001

scheduler:
  type: warmup_cosine
  warmup_epochs: 10
  warmup_start_factor: 0.1
  eta_min: 0.00001

当前启用的数据增强：

绕 Y 轴随机旋转。
小角度三轴旋转扰动。
随机等比例缩放。

其他增强策略，例如各向异性缩放、镜像、剪切、平移、jitter、dropout、局部 dropout 和点顺序打乱，已在配置文件中保留开关，可按实验需要启用。

输出文件

训练过程中会保存：

snapshots/last_model.pkl
snapshots/best_model.pkl
snapshots/epoch_XXX.pkl

其中：

last_model.pkl：最近一次训练 epoch 的 checkpoint。
best_model.pkl：验证集准确率最高的 checkpoint；如果没有验证集，则使用训练准确率。
epoch_XXX.pkl：按 snapshot_interval 周期保存的模型快照。

预测结果保存为：

result.json

格式示例：

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