PCT_Jittor

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

This repository contains the training, inference, and submission-packaging code used for a computer graphics programming assignment. The final model keeps the PCT family architecture and adds a few practical training and inference improvements for the ModelNet40 classification task.

Features

• Baseline-compatible PCT entry point in code/pct.py.
• Improved PCT recipe with unit-sphere normalization and an extra radius channel.
• SGD training with warmup cosine learning-rate scheduling.
• Test-time augmentation with point resampling, Y-axis rotations, and soft probability voting.
• Utility script for packaging predictions as the required result.json zip.

The best public score recorded during development was 0.8679.

Repository Layout

code/
  pct.py              Baseline-compatible entry point
  pct_improved.py     Improved PCT-compatible entry point
  train.py            General training entry point
  predict.py          TTA inference entry point
  make_submission.py  Package a prediction JSON as result.json in a zip
  pa3/                Dataset, model, training, inference, and utility modules
scripts/
  train_final.sh      Train the final improved PCT recipe
  predict_final.sh    Generate final predictions and submission zip
  run_final_pipeline.sh
README.md
LICENSE
requirements.txt

Generated files are intentionally excluded from version control, including data/, artifacts/, submit/, checkpoints, numpy arrays, zip files, logs, Jittor caches, and local credentials.

Environment

Jittor officially supports Linux and Windows/WSL. CUDA is recommended for training the final recipe.

Install Jittor by following the official guide:

https://cg.cs.tsinghua.edu.cn/jittor/download/

Then install the lightweight Python dependencies:

python3 -m pip install -r requirements.txt

Data

Place the ModelNet40 numpy files in data/:

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

The dataset is not stored in this repository because it is an input artifact. Keep large data files outside git and provide download instructions separately when distributing a full experiment package.

Train

Run the final training recipe:

bash scripts/train_final.sh

The expected checkpoint path is:

artifacts/models/improved_full_b32_seed2026.pkl

Predict And Package

Generate test predictions with soft TTA:

bash scripts/predict_final.sh

The prediction JSON is written to:

artifacts/results/result_final.json

The evaluation package is written to:

submit/result.zip

The zip contains one root-level file named result.json, matching the assignment evaluation format.

Final Recipe

• model: improved
• points: 1024
• batch size: 32
• seed: 2026
• learning rate: 0.003
• schedule: warmup cosine
• inference: 8 point resamples x 4 Y-axis rotations
• voting: softmax probability averaging

License

This project is released under the MIT License. See LICENSE.