|
| 1 | +--- |
| 2 | +title: TMVA SOFIE - Enhancing Keras Parser and JAX/FLAX Integration |
| 3 | +layout: gsoc_proposal |
| 4 | +project: ROOT |
| 5 | +year: 2025 |
| 6 | +organization: CERN |
| 7 | +difficulty: medium |
| 8 | +duration: 350 |
| 9 | +mentor_avail: Flexible |
| 10 | +project_mentors: |
| 11 | + |
| 12 | + organization: CERN |
| 13 | + first_name: Lorenzo |
| 14 | + last_name: Moneta |
| 15 | + is_preferred_contact: yes |
| 16 | + |
| 17 | + organization: CERN |
| 18 | + first_name: Sanjiban |
| 19 | + last_name: Sengupta |
| 20 | +--- |
| 21 | + |
| 22 | +# Description |
| 23 | +SOFIE (System for Optimized Fast Inference code Emit) is a Machine Learning Inference Engine within TMVA (Toolkit for Multivariate Data Analysis) in ROOT. SOFIE offers a parser capable of converting ML models trained in Keras, PyTorch, or ONNX format into its own Intermediate Representation, and generates C++ functions that can be easily invoked for fast inference of trained neural networks. Using the IR, SOFIE can produce C++ header files that can be seamlessly included and used in a 'plug-and-go' style. |
| 24 | + |
| 25 | +SOFIE currently supports various Machine Learning operators defined by the ONNX standards, as well as a Graph Neural Network (GNN) implementation. It supports the parsing and inference of Graph Neural Networks trained using DeepMind Graph Nets. |
| 26 | + |
| 27 | +As SOFIE continues to evolve, this project aims to: |
| 28 | +* **Enhance the Keras parser** to support models trained in the latest TensorFlow v2.18.0, which introduces NumPy 2.0 compatibility. |
| 29 | +* **Integrate JAX/FLAX support**, enabling SOFIE to generate C++ inference functions for models developed using JAX/FLAX. |
| 30 | + |
| 31 | +## Task ideas |
| 32 | +In this project, the contributor will gain experience with C++ and Python programming, TensorFlow/Keras and its storage formats for trained machine learning models, and JAX/FLAX for accelerated machine learning. They will begin by familiarizing themselves with SOFIE and its Keras parser. After researching the changes required to support the latest TensorFlow version, they will implement functionalities to ensure the successful generation of inference code for the latest Keras models. In the next phase, they will explore the JAX/FLAX library and investigate its potential integration with SOFIE. |
| 33 | + |
| 34 | +## Expected results and milestones |
| 35 | +* **Familiarization with TMVA SOFIE**: Understand the SOFIE architecture, run inference using the existing Keras parser, and analyze the current parser's capabilities. |
| 36 | +* **Researching latest TensorFlow/Keras**: Investigate the latest TensorFlow/Keras developments and assess their alignment with SOFIE. |
| 37 | +* **Improving the Keras Parser**: Implement parser enhancements to support the latest TensorFlow version and validate inference results. |
| 38 | +* **JAX/FLAX Integration**: Design and develop a parsing mechanism for JAX/FLAX models, ensuring compatibility with SOFIE’s IR and further generation of inference code. |
| 39 | + |
| 40 | +## Requirements |
| 41 | + * Proficiency in C++ and Python. |
| 42 | + * Knowledge of TensorFlow/Keras and JAX/FLAX. |
| 43 | + * Familiarity with version control systems like Git/GitHub. |
| 44 | + |
| 45 | +## Links |
| 46 | + * [ROOT Project homepage](https://root.cern/) |
| 47 | + * [ROOT Project repository](https://github.com/root-project/root) |
| 48 | + * [SOFIE Repository](https://github.com/root-project/root/tree/master/tmva/sofie) |
| 49 | + * [Keras: The high-level API for TensorFlow](https://www.tensorflow.org/guide/keras) |
| 50 | + * [JAX Documentation](https://docs.jax.dev/en/latest/) |
| 51 | + * [FLAX Documentation](https://flax.readthedocs.io/en/latest/) |
0 commit comments