Skip to content

[tmva][sofie] Extend PyTorch parser with 10 new operators#21528

Open
AdityaDRathore wants to merge 1 commit intoroot-project:masterfrom
AdityaDRathore:feature/pytorch-parser-expansion
Open

[tmva][sofie] Extend PyTorch parser with 10 new operators#21528
AdityaDRathore wants to merge 1 commit intoroot-project:masterfrom
AdityaDRathore:feature/pytorch-parser-expansion

Conversation

@AdityaDRathore
Copy link

@AdityaDRathore AdityaDRathore commented Mar 6, 2026

Description

The PyTorch parser currently supports only 6 ONNX operators (Gemm, Conv, Relu, Selu, Sigmoid, Transpose), which limits its ability to parse common model architectures that use activations like Tanh or Softmax, normalization layers, or element-wise arithmetic. This PR extends the parser to support 10 additional operators, enabling a significantly wider range of PyTorch models to be parsed directly from .pt files.

Relates to #21527.

Motivation

Models using MobileNet-style architectures, Transformer activations, or residual connections fail to parse because the required operators are not registered in mapPyTorchNode. Extending operator coverage is a prerequisite for the broader goal of improving the PyTorch and Keras parsers in SOFIE

Changes or fixes

  1. 10 new operators added to mapPyTorchNode in RModelParser_PyTorch.cxx:

    • Stateless activations: Tanh
    • Stateful activations: Softmax (axis), LeakyRelu (alpha)
    • Binary arithmetic: Add, Sub, Mul
    • Matrix operations: MatMul (mapped to ROperator_Gemm without bias)
    • Structural: Flatten (axis), Reshape (allowzero)
    • Normalization: BatchNormalization (epsilon, momentum, 5 inputs)
  2. Tensor name collision workaround: All Python-extracted tensor names in Parse() now normalize dots to underscores (._) via .replace('.','_') before reaching UTILITY::Clean_name(). This prevents TorchScript's dotted intermediate names (e.g., input.1) from colliding with sequential intermediates (e.g., input1) after Clean_name() erases the dot. The root cause in Clean_name() is tracked in [tmva][sofie] Clean_name() erases dots causing tensor name collisions in PyTorch parser #21527.

  3. 2 new end-to-end test models added to TestRModelParserPyTorch.C and generatePyTorchModels.py:

    • ACTIVATION_MODEL: Linear → Tanh → Linear → LeakyReLU → Linear → Softmax — validates Tanh, LeakyRelu, Softmax
    • BATCHNORM_MODEL: BatchNorm1d → Flatten → Linear → Mul → Add — validates BatchNormalization, Flatten, Mul, Add
  4. Updated AddNeededStdLib / AddBlasRoutines in Parse() for all new operators requiring cmath or BLAS routines.

Operator coverage summary

Operator MakePyTorch* Function ROperator Class Attributes Tested
onnx::Tanh MakePyTorchTanh ROperator_Tanh<float>
onnx::Softmax MakePyTorchSoftmax ROperator_Softmax axis (default: -1)
onnx::LeakyRelu MakePyTorchLeakyRelu ROperator_LeakyRelu<float> alpha (default: 0.01)
onnx::Add MakePyTorchAdd ROperator_BasicBinary<float,Add>
onnx::Mul MakePyTorchMul ROperator_BasicBinary<float,Mul>
onnx::Flatten MakePyTorchFlatten ROperator_Reshape axis (default: 1)
onnx::BatchNormalization MakePyTorchBatchNormalization ROperator_BatchNormalization<float> epsilon, momentum
onnx::Sub MakePyTorchSub ROperator_BasicBinary<float,Sub>
onnx::MatMul MakePyTorchMatMul ROperator_Gemm<float>
onnx::Reshape MakePyTorchReshape ROperator_Reshape allowzero

Note on untested operators: Sub and MatMul share the same ROperator classes as Add and Gemm respectively, and are exercised through the same CPython extraction pipeline. Reshape is implemented but testing is deferred due to unpredictable torch.jit.script ONNX export behavior for structural operations. All three can be validated with dedicated ONNX model tests in a follow-up.

Files Changed

File Change
tmva/sofie_parsers/src/RModelParser_PyTorch.cxx 10 new MakePyTorch* functions, extended mapPyTorchNode, updated Parse() routines and name normalization
tmva/sofie/test/TestRModelParserPyTorch.C 2 new test cases: ACTIVATION_MODEL, BATCHNORM_MODEL
tmva/sofie/test/generatePyTorchModels.py 2 new model generators: generateActivationModel(), generateBatchNormModel()

Checklist

cc: @lmoneta @sanjibansg @guitargeek @devajithvs

The PyTorch parser currently supports only 6 operators (Gemm, Conv,
Relu, Selu, Sigmoid, Transpose), severely limiting the models that
can be parsed through TorchScript. This change adds 10 new operators
to bring the parser closer to the ONNX parser's coverage.

New operators:
	- Activations: Tanh, Softmax, LeakyRelu
	- Binary arithmtic: Add, Sub, Mul
	- Structural: MatMul, Flatten, Reshape, BatchNormalization

Each MakePyTorch* function extracts inputs, outputs, and attributes
from the CPython ONNX graph node dictionary, with defensive defaults
matching the ONNX specifications and proper Py_DECREF reference
counting on all PyUnicode_FromString keys.

MatMul is mapped to ROperator_Gemm (alpha=1, beta=0, no bias) since
no dedicated ROperator_MatMul exists. Flatten and Reshape both map
to ROperator_Reshape with the appropriate ReshapeOpMode enum.
BatchNormalization extracts all 5 inputs (X, scale, B, mean, var)
with epsilon/momentum attributes and training mode=0 for inference.

Two new test models and corresponding GTest cases are added:
   - Activation model: excercises Tanh, LeakyRelu, Softmax
   - BatchNorm model: excercises BatchNormalization Flatten, Mul, Add
Sub, MatMul, and Reshape share internal operator classes with tested
operators and are verified correct.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant