Skip to content

Integrate LitePT and HEPT#474

Merged
jpata merged 26 commits into
mainfrom
jp_20260328_agg
Apr 21, 2026
Merged

Integrate LitePT and HEPT#474
jpata merged 26 commits into
mainfrom
jp_20260328_agg

Conversation

@jpata
Copy link
Copy Markdown
Owner

@jpata jpata commented Mar 30, 2026

This PR introduces major architectural updates, new model backends, and streamlined evaluation workflows.

  • New Architectures: Integrated HEPT and LitePT models; refactored MLPF to support modular attention backends (FlashAttention, SDPA).
  • Inference & Deployment: Improved ONNX export and expanded benchmarking scripts for CPU/GPU performance and memory profiling.
  • Evaluation & Validation: Added a standalone Key4HEP evaluator for ROOT files (EDM4hep) and updated the ONNX validation suite for systematic model comparisons.
  • Infrastructure & Tooling: added Hugging Face Hub integration scripts and a new detector hit visualization tool.
  • Documentation: Updated README with streamlined guides for installation, dataset handling, and model evaluation.

LitePT is a lightweight, high-performance 3D point cloud architecture. LitePT-S has 3.6 × fewer parameters, 2 × faster runtime and 2 × lower memory footprint than PTv3. Moreover, it remains faster and more memory-efficient even when scaled up to LitePT-L with a parameter count twice that of PTv3. Already the smallest variant, LitePT-S, matches or outperforms state-of-the-art point cloud backbones across a range of benchmarks.

Valid loss: HEPT noticeably worse.
Screenshot 2026-04-15 at 10 28 48

Jet match fraction: attention outperforms all.
Screenshot 2026-04-15 at 10 31 05

Jet IQR: about the same for all models.
Screenshot 2026-04-15 at 10 31 21

ONNX export check on GPU, attention has the best loss and GPU performance:
loss_vs_runtime

ONNX on CPU:
loss_vs_runtime

NB: A major caveat of this PR right now is that LitePT only supports Nvidia out of the box due to complex dependencies such as spconv-cu124. It also requires pinning torch to torch==2.9.0.

@jpata jpata changed the title Try LitePT Integrate LitePT and HEPT Apr 15, 2026
@jpata jpata merged commit 6efabd5 into main Apr 21, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant