/Tracking-ML-Exa.TrkX

Primary LanguageJupyter Notebook

Tracking with ML

Exa.TrkX Collaboration

Documentation

make_docs ci Code style: black

Welcome to repository and documentation for ML pipelines and techniques by the ExatrkX Collaboration.

Objectives

  1. To present the ExaTrkX Project pipeline for track reconstruction of TrackML and ITk data in a clear and simple way
  2. To present a set of templates, best practices and results gathered from significant trial and error, to speed up the development of others in the domain of machine learning for high energy physics. We focus on applications specific to detector physics, but many tools can be applied to other areas, and these are collected in an application-agnostic way in the Tools section.

Disclaimer:

This repository has been functional, but ugly. It is moving to an "alpha" version whereby the abstract engineering for running the pipeline has been moved to a new library TrainTrack, and keeping physics-specific research code in this repository. This transition is expected before May 2021. Please be a little patient if using before then, and if something is broken, pull first to make sure it's not already solved, then post an issue second.

Intro

To start as quickly as possible, clone the repository, Install and follow the steps in Quickstart. This will get you generating toy tracking data and running inference immediately. Many of the choices of structure will be made clear there. If you already have a particle physics problem in mind, you can apply the Template that is most suitable to your use case.

Once up and running, you may want to consider more complex ML Models. Many of these are built on other libraries (for example Pytorch Geometric).

Install

It's recommended to start a conda environment before installation:

conda create --name exatrkx-tracking python=3.8
conda activate exatrkx-tracking
pip install pip --upgrade

If you have a CUDA GPU available, load the toolkit or install it now. You should check that this is done by running nvcc --version. Then, running:

python install.py

will attempt to negotiate a path through the packages required, using nvcc --version to automatically find the correct wheels. Warning: If you are installing with cpu, this may take up to 15 minutes due to an unfortunately slow installation of Pytorch3D from source.

You should be ready for the Quickstart!

If this doesn't work, you can step through the process manually:

CPU GPU
  1. Run export CUDA=cpu

1a. Find the GPU version cuda XX.X with nvcc --version

1b. Run export CUDA=cuXXX, with XXX = 92, 101, 102, 110
  1. Install Pytorch and dependencies
    pip install --user -r requirements.txt
  1. Install local packages

pip install -e .
  1. Install CPU-optimized packages
pip install faiss-cpu
pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
  1. Install GPU-optimized packages

pip install faiss-gpu cupy-cudaXXX, with XXX

pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py3{Y}_cu{XXX}_pyt{ZZZ}/download.html

where {Y} is the minor version of Python 3.{Y}, {XXX} is as above, and {ZZZ} is the version of Pytorch {Z.ZZ}.

e.g. `py36_cu101_pyt170` is Python 3.6, Cuda 10.1, Pytorch 1.70.

Vintage Errors

A very possible error will be

OSError: libcudart.so.XX.X: cannot open shared object file: No such file or directory

This indicates a mismatch between CUDA versions. Identify the library that called the error, and ensure there are no versions of this library installed in parallel, e.g. from a previous pip --user install.