PCT_jittor

A Jittor implementation of Point Cloud Transformer (PCT) for 3D shape classification on ModelNet40.

Best validation accuracy: 81.02% (epoch 183, 200-epoch training).

Requirements

• Python 3.7+
• Jittor >= 1.3.10
• NumPy

Install Jittor:

pip install jittor

Data

Place the ModelNet40 dataset files under ./data/:

data/
├── train_points.npy    # (9843, N, 3)
├── train_labels.npy    # (9843,)
├── test_points.npy     # (2468, N, 3)
└── categories.txt      # 40 class names

Training

python pct.py --data_dir ./data

Optional arguments:

Argument	Default	Description
--data_dir	./data	Path to dataset directory
--n_points	1024	Number of input points per sample
--batch_size	32	Training batch size
--epochs	200	Total training epochs
--lr	0.001	Initial learning rate
--seed	42	Random seed

The best model is saved to best_pct_model.pkl and predictions are written to result.json.

Model Architecture

Input (B, 3, N)
  → Embedding: Conv1d(3→128) × 2 + BN + LeakyReLU
  → SA_Layer × 4  (scaled dot-product self-attention, shared Q/K projection)
  → Concat + Conv1d(512→1024) + BN + LeakyReLU
  → Global Max Pooling
  → FC(1024→512→256→40) + BN + Dropout(0.3)
Output logits (B, 40)

Total parameters: ~1.37M

Key Design Choices

• Shared Q/K projection — one Conv1d maps to both Q and K, reducing parameters
• Scaled dot-product attention — divides by √d_k to stabilize softmax gradients
• LeakyReLU (scale=0.2) — avoids dying neurons compared to ReLU
• Cosine annealing LR — η_min=1e-5, T_max=200
• Validation split — 90/10 train/val, saves the best checkpoint

Results

Metric	Value
Best validation accuracy	81.02%
Best epoch	183 / 200
Final train accuracy	79.53%
Test samples predicted	2468

Predictions are in result.json (sample index → class index mapping).

Evaluation Screenshot