PCT Jittor (ModelNet40)

Point Cloud Transformer (PCT) baseline implemented in Jittor for ModelNet40 classification. The script trains a model on the provided point clouds and produces a result.json prediction file for the test split.

Features

• PCT model with stacked self-attention blocks
• Data augmentation for point clouds
• Cosine annealing learning rate scheduler
• Training and test prediction in one script

Requirements

• Python 3.10+ recommended
• Linux (GPU optional)
• Dependencies in requirements.txt

Install dependencies:

pip install -r requirements.txt

Quick Start

Train and generate predictions:

python pct.py

Common options:

python pct.py --epochs 200 --batch_size 32 --lr 0.002 --n_points 1024

120 epoches is enough to get 80% acc.

Outputs:

• pct_model.pkl: saved model checkpoint
• result.json: test predictions (key: sample id, value: class id)

Dataset

Expected files under data/:

• train_points.npy (N, 2048, 3)
• train_labels.npy (N,)
• test_points.npy (M, 2048, 3)
• categories.txt (class names)

Docker

CPU image:

docker build -f Dockerfile.cpu -t pct-jittor:cpu .
docker run --rm -it -v "$PWD":/workspace pct-jittor:cpu

GPU image (requires NVIDIA runtime):

docker build -f Dockerfile.gpu -t pct-jittor:gpu .
docker run --rm -it --gpus all -v "$PWD":/workspace pct-jittor:gpu

Notes

• Default is CUDA enabled in pct.py with jt.flags.use_cuda = 1. Set it to 0 if you want CPU only.
• The provided test.py is a minimal Jittor sanity check.