PCT_jittor

A Jittor implementation of point cloud classification models for the ModelNet40 assignment. The repository includes the assignment PCT baseline and an optional DGCNN model, both trained and evaluated with Jittor.

Features

• ModelNet40 point cloud classification with preprocessed .npy files.
• PCT baseline model implemented in Jittor.
• DGCNN model implemented in Jittor for stronger performance.
• Data augmentation: random rotation, random scaling, and jitter.
• Test-time augmentation (TTA) for prediction.
• Probability ensembling utility for combining multiple submissions.

Repository Layout

.
|-- pct.py              # Training and inference entry point
|-- ensemble_probs.py   # Average multiple probability files into one result
|-- environment.yml     # Conda environment
|-- requirements.txt    # Pip dependencies
|-- README.md
`-- .gitignore

Large generated files are intentionally not tracked, including datasets, checkpoints, probability arrays, result JSON files, and zip submissions.

Environment

Create the conda environment:

conda env create -f environment.yml
conda activate jitu

Or install dependencies with pip:

pip install -r requirements.txt

Tested environment:

• Python 3.9
• Jittor 1.3.10.0
• NumPy 1.26.4
• CUDA-enabled Jittor runtime for GPU training

Data

Download the official preprocessed ModelNet40 data:

https://cloud.tsinghua.edu.cn/f/f003de5a2e914d1e9e0e/?dl=1

Extract the files into data/:

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

The dataset is ignored by git because it is large.

Usage

Quick smoke test

python pct.py \
  --data_dir ./data \
  --epochs 1 \
  --batch_size 2 \
  --n_points 128 \
  --max_train_batches 2

Train PCT

python pct.py \
  --model pct \
  --data_dir ./data \
  --epochs 200 \
  --batch_size 32 \
  --n_points 2048 \
  --lr 0.01 \
  --seed 42 \
  --aug_scale \
  --aug_jitter \
  --tta_views 8 \
  --save pct_2048_aug_s42.pkl \
  --result result_2048_aug_s42_tta8.json \
  --probs probs_2048_aug_s42_tta8.npy

Train DGCNN

python pct.py \
  --model dgcnn \
  --data_dir ./data \
  --epochs 250 \
  --batch_size 32 \
  --n_points 1024 \
  --k 20 \
  --emb_dims 1024 \
  --dropout 0.5 \
  --lr 0.1 \
  --weight_decay 1e-4 \
  --seed 42 \
  --aug_scale \
  --aug_jitter \
  --tta_views 16 \
  --save dgcnn_1024_s42.pkl \
  --result result_dgcnn_1024_s42_tta16.json \
  --probs probs_dgcnn_1024_s42_tta16.npy

Evaluate from a checkpoint

python pct.py \
  --model dgcnn \
  --data_dir ./data \
  --eval_only \
  --load_model dgcnn_1024_s42.pkl \
  --batch_size 32 \
  --n_points 1024 \
  --k 20 \
  --emb_dims 1024 \
  --dropout 0.5 \
  --tta_views 16 \
  --result result.json

Ensemble probabilities

python ensemble_probs.py \
  --probs probs_a.npy probs_b.npy \
  --result result_ensemble.json

Submission

The script writes a JSON mapping from test sample id to class id. Compress the result JSON file before submitting to the course platform:

zip result.zip result.json

Notes

• Checkpoints (*.pkl), datasets, probability files, and submission outputs are ignored by git.
• The best local submission in this development run reached 89.75 on the course platform with a DGCNN-based configuration.

References

• Point Cloud Transformer (PCT), Guo et al.
• Dynamic Graph CNN for Learning on Point Clouds, Wang et al.
• Jittor: https://github.com/Jittor/jittor
• README guide: https://chinese.freecodecamp.org/news/how-to-write-a-good-readme-file/
• gitignore documentation: https://git-scm.com/docs/gitignore