PCT_GYJ

PCT_GYJ: Point Cloud Transformer with Jittor

一个基于 Jittor 框架实现的 Point Cloud Transformer (PCT) 点云分类项目，用于 PA3 ModelNet40 分类任务。

项目信息

项目名称：PCT_GYJ
项目简介：A Jittor implementation of Point Cloud Transformer (PCT) for ModelNet40 classification.
使用框架：Jittor
任务类型：点云分类
数据集：ModelNet40

本项目实现了一个用于三维点云分类的 PCT 模型。代码包含数据读取、点云预处理、模型训练、验证、模型保存以及测试集预测结果生成。

文件结构

.
├── README.md          # 项目说明文档
├── .gitignore         # Git 忽略规则
├── pct.py             # PCT 模型、训练、验证和预测代码
└── PA3_README.pdf     # 作业说明文件

训练过程中会生成以下文件，这些文件不需要上传到 Gitlink：

pct_model.pkl          # 最优验证模型权重
result.json            # 测试集预测结果
train.log              # 训练日志

数据集

数据集文件较大，不放入 Git 仓库。请从下面链接下载数据集：

https://cloud.tsinghua.edu.cn/d/33239a0006304dc9b239/

下载后将数据文件放到项目根目录下的 data/ 文件夹中，目录结构如下：

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

其中 train_points.npy 和 train_labels.npy 用于训练和验证，test_points.npy 用于生成最终提交的测试集预测结果。

环境要求

推荐使用 Docker 或 Conda 配置 Jittor 环境，避免污染服务器原有环境。

基本要求：

Python >= 3.7
g++ >= 5.4.0
Jittor
CUDA 和 cuDNN，推荐用于 GPU 训练

安装 Jittor：

python -m pip install jittor

如果使用 Docker，可以参考 Jittor 官方镜像或带 CUDA 的基础镜像运行本项目。

训练

在项目根目录下运行：

python pct.py --data_dir ./data --epochs 200 --batch_size 16 --n_points 1024 --num_workers 4 --votes 6

参数说明：

--data_dir：数据集目录
--epochs：训练轮数
--batch_size：批大小
--n_points：每个点云采样的点数
--num_workers：数据加载线程数
--votes：测试时投票次数

训练完成后会保存最优模型 pct_model.pkl，并生成测试集预测文件 result.json。

GPU Docker 示例

下面命令会在 Docker 中运行训练，并将 Jittor 缓存放到项目目录外：

docker run --rm --gpus all --ipc=host \
  -e JITTOR_HOME=/jittor_home \
  -e cache_name=pa3_train \
  -v /tmp/pa3_jittor_home:/jittor_home \
  -v $(pwd):/workspace \
  -w /workspace \
  nvcr.io/nvidia/pytorch:24.07-py3 \
  bash -lc "python -m pip install -q jittor && python pct.py --data_dir ./data --epochs 200 --batch_size 16 --n_points 1024 --num_workers 4 --votes 6"

查看 GPU 使用情况：

nvidia-smi

预测

如果已经有训练好的模型，可以直接生成测试集预测结果：

python pct.py --data_dir ./data --predict_only --resume pct_model.pkl

输出文件为：

result.json

模型简介

本项目中的 PCT 模型主要包含：

点云输入归一化和随机采样
点云旋转、缩放和平移增强
点特征提取层
多层 self-attention 模块
全局特征聚合
分类器输出 40 个 ModelNet40 类别

训练策略包括 SGD 优化、动量、余弦学习率衰减、验证集选择最优模型以及测试时旋转投票。

结果

本项目训练完成后会在验证集上保存表现最好的模型，并使用该模型对测试集生成 result.json。该文件可用于作业平台提交。

Gitignore 说明

.gitignore 已忽略以下内容：

数据集目录 data/
NumPy 数据文件 *.npy
Jittor 编译缓存 .jittor_home/
Python 缓存 __pycache__/
训练日志 *.log
模型权重 pct_model.pkl
预测结果 result.json

这些文件通常由程序自动生成或体积较大，不适合上传到 Gitlink 仓库。