This repo is to reproduce the GlueX tracking algorithm with PyTorch, which originally implemented with TensorFlow Keras here. It aims to a future integration with phasm.

To the best of my knowledge, mimic everything in the original Keras notebook, including the same:

• Shuffled training dataset;
• Batch size, epochs, NN size, loss function, optimizer, clip value;
• Learning rate scheduler.

• utils.py: define the NN structure and wrap the datasets into batched PyTorch dataloaders.
• LSTM_training.py: train the NN with the whole training dataset with 100 epochs. Save the trained model into a TorchScript.
• validation_processing.py: load the model from the TorchScript. Validate the model with a validation dataset of (661644, 6).
• submit-training-job.slurm: a one-step slurm script to run jobs on farm.

Based on my own experience, the bare-metal python3.9+pip3+cudnn8.6 installation would always fail on JLab ifarm A100 GPU nodes because of mismatch cudnn/torch versions. This is solved by installing the latest pytorch (as of Nov-28-2022) via conda virtual environments as guided here. A conda environment file is provided to show my environment configurations.

# install pytorch via conda
conda install pytorch torchvision torchaudio pytorch-cuda=11.6 -c pytorch -c nvidia
# create conda env from yml file
conda env create -f environment.yml

Table: the LSTM network definition, where batch_size=1256 and seq_len=7.

The parameter counts of the layers are taken from the original Keras model.summary().

Download the training dataset as below.

wget https://halldweb.jlab.org/talks/ML_lunch/Sep2019/MLchallenge2_training.csv
mv MLchallenge2_training.csv train_data.csv

Compared to the dataset at the time of executing the Keras notebook, the new dataset is about 38.5% larger (2646573 v.s. 1910698).

After sequencing, the dimension of the whole training dataset (as on 10/20/2022) is (2646573, 7, 6), with each epoch containing ~2108 batches. We train 100 epochs in total.

Table: results after 100 training epochs

The code is tested on a single ifarm TitanRTX/T4/A100 GPU. Results are available at:

• ./res/training-loss: images of the losses along the training process.
• ./res/job-logs: the detailed job logs. An example of how losses are changed along the epochs, batches and time is here.
• ./res/evaluation: images of the evaluation results. This is a comparison between the evaluation errors with Epochs=1 and Epochs=100.

• Keras APIs: https://keras.io/api/
• PyTorch documentation: https://pytorch.org/docs/stable/index.html
• PyTorch tutorials
  • Saving and Loading Models
  • From a LSTM cell to a Multilayer LSTM Network with PyTorch
  • Loading a TorchScript model in C++

Last updated on 02/01/2023 by [email protected]