PCT_Jittor

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

About The Project

This repository contains a course assignment implementation for 3D point cloud classification on the ModelNet40 dataset. The project uses the Jittor deep learning framework and follows the Point Cloud Transformer (PCT) architecture for 40-class shape classification.

After training, the script runs prediction on the test set and exports the results to result.json.

Features

• Point cloud classification with PCT
• Jittor-based training and inference pipeline
• Random point sampling for each shape
• Training-time point cloud augmentation
• Cosine annealing learning rate schedule
• Export of test predictions to result.json

Built With

• Python 3.9+
• Python
• Jittor
• NumPy

Project Structure

PCT_Jittor/
├── LICENSE
├── pct.py
├── README.md
├── requirements.txt
├── scripts/
├── outputs/
└── data/

Getting Started

Prerequisites

Install the required Python packages:

pip install -r requirements.txt

If Jittor requires additional compiler or CUDA setup on your machine, refer to the official installation guide:

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

Dataset Preparation

Place the dataset files under data/:

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

Expected shapes:

• train_points.npy: (N, 2048, 3)
• train_labels.npy: (N,)
• test_points.npy: (M, 2048, 3)

Large dataset files are not included in this repository.

The dataset used in this assignment is provided by the course platform and can be downloaded from:

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

The data root is configured through the command-line argument --data_dir.

Usage

Run training and prediction with default settings:

python pct.py --device auto

Example with common arguments:

python pct.py \
  --data_dir ./data \
  --n_points 1024 \
  --batch_size 32 \
  --epochs 200 \
  --lr 0.01 \
  --seed 42 \
  --device auto

Main Arguments

• --data_dir: dataset directory
• --n_points: number of sampled points per shape
• --batch_size: batch size
• --epochs: number of training epochs
• --lr: learning rate
• --device: auto, cpu, or cuda
• --save_path: model output path
• --result_path: prediction output path

The script controls random seeds through --seed and sets NumPy, Python, and Jittor seeds inside the entry script.

Training

Run training with the provided shell script:

bash scripts/train.sh

You can also run the training command directly:

python pct.py \
  --data_dir ./data \
  --n_points 1024 \
  --batch_size 32 \
  --epochs 200 \
  --lr 0.01 \
  --seed 42 \
  --device auto

Evaluation / Inference

This baseline script performs prediction on the test set immediately after training and writes the output to result.json.

Example command:

python pct.py \
  --data_dir ./data \
  --save_path pct_model.pkl \
  --result_path result.json \
  --device auto

Output Files

After training, the script generates:

• pct_model.pkl: trained model weights
• result.json: predicted labels for the test set

Example format of result.json:

{
  "0": 12,
  "1": 7,
  "2": 31
}

Each key is the test sample index as a string, and each value is the predicted class id.

Training Details

• Optimizer: SGD
• Momentum: 0.9
• Weight decay: 1e-4
• Learning rate schedule: cosine annealing
• Loss function: cross-entropy loss

Result Notes

• The local script exports test predictions as result.json, which is the file required for submission.
• The test set provided in this assignment does not include ground-truth labels, so the script does not report a final test accuracy metric locally.
• Final leaderboard or platform results may differ from local observations because the online system evaluates the hidden test labels after submission.

Notes

• The training set includes labels, while the test set only contains point clouds.
• The script predicts the labels of test samples and saves them to result.json.
• Generated model files, prediction files, dataset files, and cache files are excluded by .gitignore.