Final Evaluation Report for GSoC 2025

Details

Project Description

The SOFIE (System for Optimized Fast Inference Code Emit) project is an initiative within the TMVA (Toolkit for Multivariate Data Analysis) framework in ROOT, which aims to enhance the efficiency and speed of inference for machine learning models. SOFIE converts ML models trained in different frameworks, such as ONNX, PyTorch, and TensorFlow, into an Intermediate Representation (IR). This IR allows SOFIE to generate optimized C++ functions for fast and effective inference of neural networks and subsequently convert them into C++ header files, which can be used in a plug-and-go style for inference.

SOFIE’s workflow

To reduce the overhead of using multiple frameworks for inference, SOFIE generates unified inference code for models trained with different frameworks.

SOFIE has mainly two components: a Parser and an inference code generator.

SOFIE currently supports parsing mechanisms for ML models built with frameworks such as ONNX, PyTorch, and TensorFlow.

About SOFIE’s Keras Parser

Currently, SOFIE’s existing Keras parser is written in C++ and is quite old. Although it’s written in C++, the actual parsing logic is implemented in Python. Additionally, it lacks support for parsing layers such as Pooling and LayerNormalization.

Project Objectives

• Rewrite the Keras model parser in Python, replacing earlier C++ logic to improve modular design and flexibility, and simplify future extensions
• Extend parser functionality to support Pooling and LayerNormalization layers
• Enable support for Keras 3, while preserving support for Keras 2.x models, ensuring full backward compatibility
• Add support for both types of models, i.e., models built using Keras’ Functional as well as Sequential API.
• Design comprehensive unit tests for the parser to guarantee robustness and correctness

Work Accomplished

Since SOFIE’s operators are written entirely in C++, we had to leverage ROOT’s Pythonization functionality, which essentially allows us to use SOFIE’s C++ objects in a Pythonic interface. The overall structure of the parser is very similar to the previous one. The sequence of operations is as follows:

1. Load the Keras model

2. Instantiate the RModel class

3. Iterate over the individual layers and extract the required information

To create the RModel object, we had to extract layer-specific information such as Layer name, its type (Dense, Convolution, etc.), and its input and output layer names. In case of Keras 2.x and models built using Functional API, the output name of the current layer is the same as the input name for the next layer. But in the case of Keras 3, particularly with models built using the Sequential API, this changes, and the input and output names are no longer consistent (issue link). So, we used a custom iterator that would iterate over each of the layers and replace the suffix of the input and output names to be perfectly consistent.

Then we had to extract the weight’s name and some more operator-specific information, such as in the case of Convolutional and Pooling layers, ONNX only supports channels_first data format. Whereas Keras supports both, i.e., channels_first and channels_last. After extracting the layer information for a particular layer, we add it to the rmodel object.

4. Adding a layer to the rmodel object

For all the other operators, the process of adding layer operators to the rmodel object is quite easy. But for Convolutional and Pooling layers, it’s a bit different. For these layers, if the data format is channels_last, we have to perform a transpose before and after adding the layer to the rmodel object.

5. Operator-specific functions

To add the layer operators, we create the layer operator in layer-specific functions and return it. For this, we make use of the extracted layer information from step 3.

6. Extract the model’s weights

7. Adding input and output names of the Keras model to the rmodel object

While adding the input and output names of the Keras model, we need to make sure that we use the new layer iterator; otherwise, the layer names would have been inconsistent again.

How to make sure the parser has backwards compatibility with Keras version 2.x?

Along with parsing support for models trained with Keras 3, we also needed backward compatibility with Keras 2.x. Since Keras 3 introduced significant changes in attribute and layer names and storage formats, we conducted research on the updated versions. For example, the weight names in Keras 3 are no longer unique. Assume a model has 2 dense layers. Previously, with Keras 2.x, the layer weight names would have been:

dense/kernel:0
dense/bias:0
dense_1/kernel:0
dense_1/bias:0

But with Keras 3, the layer weight names are like this:

kernel
bias
kernel
bias

To remove the ambiguity, we used weight paths instead of weight names.

After these steps, the parser was in good shape and can now be used to parse these layers:

• Add
• AveragePool2D channels first
• AveragePool2D channels last
• BatchNormalization
• Concat
• Conv2D channels first
• Conv2D channels last
• Conv2D padding same
• Conv2D padding valid
• Dense
• Elu
• Flatten
• GlobalAveragePool2D channels first
• GlobalAveragePool2D channels last
• LayerNormalization
• LeakyReLU
• MaxPool2D channels first
• MaxPool2D channels last
• Multiply
• Permute
• Relu
• Reshape
• Selu
• Sigmoid
• Softmax
• Subtract
• Swish
• Tanh

Along the way, we also fixed a few minor bugs within SOFIE’s ROperator header files.

About JAX/FLAX parser

Initially, we aimed for JAX/FLAX Integration within SOFIE by researching models built using its nnx and linen API, but after a careful discussion with the project mentors, we came to the conclusion of focusing on the Keras parser itself, by adding support to parse more layers and by writing unit tests for the same.

Writing Unit Tests

Along with verifying the parsing capability to parse all the supported layers, we also needed to verify the correctness of the generated code. For this, we created two functions:

1. To generate and test the inference code

This function takes the file path of a model built with Keras. Then, it passes the same to the parser. After the parser returns the rmodel object, we generate the inference code. Now, we need to verify the correctness of the generated header file. For this, we need to pass a sample input to both the generated header file and the Keras model. To avoid hardcoding the shape of the input for each and every model, we extract the input shapes from the Keras model and using it we create the sample input. After this, we pass the sample input and get the resultant output tensor. Since SOFIE always flattens the output tensor before returning it, we also check the output tensor shape from both Keras and SOFIE.

2. Validate the accuracy of the result

To validate the inference result from SOFIE, we compare the element-wise values in the output tensor of Keras and SOFIE and make sure that the difference between the results is within a specified tolerance.

Unit tests

To write the unit tests, we used the unittest module within Python as it allows parametrizing those with minimum code repetition. There are two different sets of tests:

1. For models built using Keras’ Functional API
2. For models built using Keras’ Sequential API

Within these, there are operator-specific tests that are invoked whenever a sub-test is called. While running the unit tests for both types of models, i.e., Functional and Sequential, temporary directories are created and torn down as soon as both are finished running.

Pull request status

Pull Request	PR Number	Status
New Keras Parser	#19692

Challenges faced and Learning Outcomes

• Faced difficulty while setting up the ROOT project with SOFIE enabled due to missing dependencies and incompatible versions of packages
• Navigated through SOFIE’s complex code base
• Got hands-on experience with Keras, its Functional and Sequential API, and the overall structure of its models
• Improved skills in reading documentation and solving bugs independently
• Learned how to write concise and modular unit tests

Future work

In the future, I would love to continue my contributions to the SOFIE codebase beyond the GSoC period. My current focus is on adding support for parsing the Conv2DTranspose, Dropout, and Recurrent layers.

Conclusion

I am thankful for my project mentors, Sanjiban Sengupta and Dr. Lorenzo Moneta, for their kind guidance, which made my learning experience enriching and rewarding. They guided me whenever I felt any difficulty. I could ask them the silliest doubt, and they would still answer it happily. I am fortunate to be part of such a wonderful project and contribute to CERN-HSF this summer. I look forward to contributing to CERN-HSF beyond my GSoC project completion.

Lastly, I would like to thank my seniors from the Open-Source community at my university, Project-X, for introducing me to GSoC and helping me in the pre-GSoC period.

Thanks and Regards

Prasanna Kasar