Integrating ML models through PyTorch

[anthoak13]

Goal

To integrate the classification of events using PyTorch models. The training will be handled in python, it is just the classification that needs to be injected into the ATTPCROOT analysis stream. The generation of the tensor input should be done fully within the ATTPCROOT framework (because when classifying we are not calling any external python code)

Design Requirements

• Trained model from PyTorch, serialized into a file with TorchScript
• Ability to convert ATTPCROOT data (AtEvent/AtRawEvent) into a tensor input.
• Ability to save tensor input to an HDF5 file for training
• Ability to use trained model and tensor input to classify an event

Proposed design

There are two main additional tasks needed in the code.

1. Task to generate tensor from data and save to disk (in AtReconstruction?)
   • Proposed name: AtTensorTask
2. Task to classify events based on TouchScript model (in AtReconstruction?)
   • Proposed name: AtClassifierTask

In addition there will also be a namespace containing classes for generating input tensors from the ATTPCROOT data. These classes will be added to the LibTorch namespace torch. Our methods for generating tensors should follow the conventions established by LibTorch

AtTensorTask

This will be a wrapper task (similar to AtPSATask) where the bulk of the logic of generating the tensors lives in other classes that are more removed from ROOT.

• Input
  • std::unique_ptr<AtTensorConverter> (abstract base class) or maybe an enum indicating which tensor generation method to use.
  • Filename to save tensor to
• Output
  • HDF5 file filled with the tensors to use as an input for all events in the run (expect simulated)

AtClassifierTask

This task is responsible for classifying events based on a PyTorch model.

• Input
  • std::unique_ptr<AtTensorConverter> (abstract base class) or maybe an enum indicating which tensor generation method to use.
  • File path to TorchScript model
• Output
  • AtClassifiedEvent (Derived AtEvent with a field storing classification info)

Tensor generation

This section is by far the most nebulous. I haven't worked with PyTorch at all so am basing this mostly on a skim of the documentation. I think this will cohere into something reasonable as development continues. My initial idea was to have an abstract base class for the generation of tensors. Each experiment would then have their own, that made heavy use of private inheritance (or containment) for the creation of common components of a picture (pad plane for different detectors, image of the mesh signal, indication of total charge in an event etc). It is then simply the job of the experiment-specific class to combine all of these individual components into a single tensor. While testing the implementation of the code it would also be useful to have the ability to read a tensor from a h5 file (for example comparing with the full python implementation used by the GADGET group to make sure we can reproduce the same results).

It looks like the existing tensor generation methods of LibTorch follow a different convention using free functions grouped into a namespace rather than inheritance through classes. Following the conventions established by LibTorch we will also have a header file that defines a number of free functions in the torch namespace that operate as the factory for each way of generating tensors (think something similar to the use of AtEstimators in SampleConsensus where they exist as free functions in a namespace). It is likely that each of experiment or physics "problem" will end up implementing their own free function in this namespace. There should be a number of helper functions for generating these tensors for common tasks, which can then be spiced together. This maps almost directly onto my first idea, but you loose the ability to store a "state" which I don't think would be a huge deal. It depends on implementation choices.

Dependency changes

There are CPU-only and CUDA versions of LibTorch. I would like to support both, but I think the ability to do that well will depend heavily on how good there CMake is. Ideally, if we found both versions we would use CUDA-enabled if CUDA is available and fall back to CPU only if necessary.