diff --git a/tmva/sofie/README.md b/tmva/sofie/README.md index 144944e8f58e7..3a992fb32e888 100644 --- a/tmva/sofie/README.md +++ b/tmva/sofie/README.md @@ -89,6 +89,7 @@ std::vector supportedOperators = parser.GetRegisteredOperators(); - [x] Abs - [x] Add +- [x] ArgMax - [x] AveragePool - [x] BatchNormalization - [x] Cast diff --git a/tmva/sofie/inc/TMVA/ROperator_ArgMax.hxx b/tmva/sofie/inc/TMVA/ROperator_ArgMax.hxx new file mode 100644 index 0000000000000..199150dc90def --- /dev/null +++ b/tmva/sofie/inc/TMVA/ROperator_ArgMax.hxx @@ -0,0 +1,111 @@ +#ifndef TMVA_SOFIE_ROPERATOR_ARGMAX +#define TMVA_SOFIE_ROPERATOR_ARGMAX + +#include "TMVA/SOFIE_common.hxx" +#include "TMVA/ROperator.hxx" +#include "TMVA/RModel.hxx" + +#include +#include +#include +#include + +namespace TMVA { +namespace Experimental { +namespace SOFIE { + +class ROperator_ArgMax final : public ROperator { +private: + /* Attributes */ + int64_t fAttrAxis = 0; + int fKeepdims = 1; + int fSelectLastIndex = 0; + + /* Tensor names */ + std::string fNX; // input + std::string fNY; // output + + /* Tensor shapes */ + std::vector fShapeX; + std::vector fShapeY; + +public: + ROperator_ArgMax() {} + ROperator_ArgMax(int64_t axis, int keepdims, int selectLastIndex, std::string nameX, std::string nameY) + : fAttrAxis(axis), + fKeepdims(keepdims), + fSelectLastIndex(selectLastIndex), + fNX(UTILITY::Clean_name(nameX)), + fNY(UTILITY::Clean_name(nameY)) + { + fInputTensorNames = {fNX}; + fOutputTensorNames = {fNY}; + } + + void Initialize(RModel &model) override + { + if (!model.CheckIfTensorAlreadyExist(fNX)) { + throw std::runtime_error("TMVA SOFIE ArgMax Op Input Tensor " + fNX + " is not found in model"); + } + fShapeX = model.GetDimTensorShape(fNX); + int64_t rank = static_cast(fShapeX.size()); + + // negative axis counts from the end (e.g. -1 is the last dimension) + if (fAttrAxis < 0) { + fAttrAxis += rank; + } + // valid range is [-rank, rank-1], i.e. [0, rank-1] after normalization + if (fAttrAxis < 0 || fAttrAxis >= rank) { + throw std::runtime_error("TMVA SOFIE ArgMax Op - invalid axis value " + std::to_string(fAttrAxis) + + " for input tensor " + fNX + " of rank " + std::to_string(rank)); + } + + // compute output shape: reduced axis becomes 1 (keepdims) or is removed + fShapeY = fShapeX; + if (fKeepdims != 0) { + fShapeY[fAttrAxis] = Dim{1}; + } else { + fShapeY.erase(fShapeY.begin() + fAttrAxis); + } + // output contains indices, so its type is always INT64 regardless of the input type + model.AddIntermediateTensor(fNY, ETensorType::INT64, fShapeY); + } + + std::string Generate(std::string opName) override + { + // split the input shape around the axis: [outer..., axisDim, inner...] + std::string outer = ConvertDimShapeToLength(std::vector(fShapeX.begin(), fShapeX.begin() + fAttrAxis)); + std::string axisDim = fShapeX[fAttrAxis].GetVal(); + std::string inner = ConvertDimShapeToLength(std::vector(fShapeX.begin() + fAttrAxis + 1, fShapeX.end())); + + // element (o, k, i) of the input is at o * axisDim * inner + k * inner + i + std::string stride = "(" + axisDim + ") * (" + inner + ")"; + // with '>' ties keep the first index, with '>=' they move to the last one + std::string cmp = (fSelectLastIndex != 0) ? ">=" : ">"; + + std::stringstream out; + out << "\n//------ ArgMax " << opName << "\n"; + out << SP << "for (size_t o = 0; o < " << outer << "; o++) {\n"; + out << SP << SP << "for (size_t i = 0; i < " << inner << "; i++) {\n"; + out << SP << SP << SP << "size_t base = o * " << stride << " + i;\n"; + out << SP << SP << SP << "auto greatest = tensor_" << fNX << "[base];\n"; + out << SP << SP << SP << "int64_t bestIndex = 0;\n"; + out << SP << SP << SP << "for (size_t k = 1; k < " << axisDim << "; k++) {\n"; + out << SP << SP << SP << SP << "auto value = tensor_" << fNX << "[base + k * (" << inner << ")];\n"; + out << SP << SP << SP << SP << "if (value " << cmp << " greatest) {\n"; + out << SP << SP << SP << SP << SP << "greatest = value;\n"; + out << SP << SP << SP << SP << SP << "bestIndex = static_cast(k);\n"; + out << SP << SP << SP << SP << "}\n"; + out << SP << SP << SP << "}\n"; + out << SP << SP << SP << "tensor_" << fNY << "[o * (" << inner << ") + i] = bestIndex;\n"; + out << SP << SP << "}\n"; + out << SP << "}\n"; + return out.str(); + } +}; + +} // namespace SOFIE +} // namespace Experimental +} // namespace TMVA + +#endif // TMVA_SOFIE_ROPERATOR_ARGMAX diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx index b65c98ae29cc4..f39859ca9fba2 100644 --- a/tmva/sofie/test/TestCustomModelsFromONNX.cxx +++ b/tmva/sofie/test/TestCustomModelsFromONNX.cxx @@ -615,6 +615,38 @@ TEST(ONNX, ReduceMin) expectNear(output, correct_output, DEFAULT_TOLERANCE); } +TEST(ONNX, ArgMax) +{ + // default attributes: axis=0, keepdims=1 on a [2,3] tensor + std::vector input({1, 9, 3, 7, 2, 8}); + std::vector correct_output({1, 0, 1}); + + ASSERT_INCLUDE_AND_RUN(std::vector, "ArgMax", input); + + expectEqual(output, correct_output); +} + +TEST(ONNX, ArgMax_Axis1_NoKeepdims) +{ + std::vector input({1, 9, 3, 7, 2, 8}); + std::vector correct_output({1, 2}); + + ASSERT_INCLUDE_AND_RUN(std::vector, "ArgMax_Axis1_NoKeepdims", input); + + expectEqual(output, correct_output); +} + +TEST(ONNX, ArgMax_NegAxis_SelectLast) +{ + // axis=-1 with ties: select_last_index=1 picks the last maximum + std::vector input({5, 5, 1, 2, 7, 7}); + std::vector correct_output({1, 2}); + + ASSERT_INCLUDE_AND_RUN(std::vector, "ArgMax_NegAxis_SelectLast", input); + + expectEqual(output, correct_output); +} + // Elu on a tensor whose first dimension is only known at run time. TEST(ONNX, EluDynShape) { diff --git a/tmva/sofie/test/generate_input_models.py b/tmva/sofie/test/generate_input_models.py index 0d649c4cd0c1f..5dd39ba342e08 100644 --- a/tmva/sofie/test/generate_input_models.py +++ b/tmva/sofie/test/generate_input_models.py @@ -245,6 +245,39 @@ def make_AddBroadcast7(): return _model(graph, opset=17, ir_version=8) +def _argmax_model(name, out_shape, **attrs): + """ArgMax on a [2, 3] float tensor with the given attributes.""" + nodes = [ + helper.make_node("ArgMax", ["input"], ["output"], **attrs), + ] + graph = helper.make_graph( + nodes, + name, + inputs=[ + _vi("input", FLOAT, [2, 3]), + ], + outputs=[ + _vi("output", INT64, out_shape), + ], + ) + return _model(graph, opset=13, ir_version=8) + + +def make_ArgMax(): + """Ops: ArgMax (default attributes: axis=0, keepdims=1)""" + return _argmax_model("argmax_test", [1, 3]) + + +def make_ArgMax_Axis1_NoKeepdims(): + """Ops: ArgMax""" + return _argmax_model("argmax_axis1_test", [2], axis=1, keepdims=0) + + +def make_ArgMax_NegAxis_SelectLast(): + """Ops: ArgMax""" + return _argmax_model("argmax_select_last_test", [2, 1], axis=-1, keepdims=1, select_last_index=1) + + def make_AveragePool1d_CeilMode_Overhang(): """Ops: AveragePool""" nodes = [ @@ -6164,6 +6197,9 @@ def make_AliasDynShapeAcrossNewTensor(): "AddBroadcast5": make_AddBroadcast5, "AddBroadcast6": make_AddBroadcast6, "AddBroadcast7": make_AddBroadcast7, + "ArgMax": make_ArgMax, + "ArgMax_Axis1_NoKeepdims": make_ArgMax_Axis1_NoKeepdims, + "ArgMax_NegAxis_SelectLast": make_ArgMax_NegAxis_SelectLast, "Asinh": make_Asinh, "Atanh": make_Atanh, "AveragePool1d_CeilMode": make_AveragePool1d_CeilMode, diff --git a/tmva/sofie_parsers/CMakeLists.txt b/tmva/sofie_parsers/CMakeLists.txt index 12c0ccba85612..2402b725dc50e 100644 --- a/tmva/sofie_parsers/CMakeLists.txt +++ b/tmva/sofie_parsers/CMakeLists.txt @@ -78,6 +78,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofieParser src/ParseGelu.cxx src/ParseHardSigmoid.cxx src/ParseHardSwish.cxx + src/ParseArgMax.cxx DEPENDENCIES ROOTTMVASofie ) diff --git a/tmva/sofie_parsers/src/ParseArgMax.cxx b/tmva/sofie_parsers/src/ParseArgMax.cxx new file mode 100644 index 0000000000000..a8f224e87a64b --- /dev/null +++ b/tmva/sofie_parsers/src/ParseArgMax.cxx @@ -0,0 +1,44 @@ +#include "TMVA/RModelParser_ONNX.hxx" +#include "TMVA/ROperator_ArgMax.hxx" +#include "onnx.hxx" + +namespace TMVA { +namespace Experimental { +namespace SOFIE { + +ParserFuncSignature ParseArgMax = [](RModelParser_ONNX &parser, const onnx::NodeProto &nodeproto) { + auto input_name = nodeproto.input(0); + if (!parser.IsRegisteredTensorType(input_name)) { + throw std::runtime_error("TMVA::SOFIE ONNX Parser ArgMax op has input tensor " + input_name + + " but its type is not yet registered"); + } + + // default values as in the ONNX specification + int64_t attr_axis = 0; + int attr_keepdims = 1; + int attr_select_last_index = 0; + for (int i = 0; i < nodeproto.attribute_size(); i++) { + const std::string &attribute_name = nodeproto.attribute(i).name(); + if (attribute_name == "axis") + attr_axis = nodeproto.attribute(i).i(); + else if (attribute_name == "keepdims") + attr_keepdims = nodeproto.attribute(i).i(); + else if (attribute_name == "select_last_index") + attr_select_last_index = nodeproto.attribute(i).i(); + } + + std::string output_name = nodeproto.output(0); + std::unique_ptr op( + new ROperator_ArgMax(attr_axis, attr_keepdims, attr_select_last_index, input_name, output_name)); + + // the output contains indices, so it is always INT64 + if (!parser.IsRegisteredTensorType(output_name)) { + parser.RegisterTensorType(output_name, ETensorType::INT64); + } + + return op; +}; + +} // namespace SOFIE +} // namespace Experimental +} // namespace TMVA diff --git a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx index 3724c72ca8ed2..7ec8dbfbfd1d9 100644 --- a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx +++ b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx @@ -84,6 +84,7 @@ extern ParserFuncSignature ParseScatterElements; extern ParserFuncSignature ParseScatterND; extern ParserFuncSignature ParseNonZero; extern ParserFuncSignature ParseClip; +extern ParserFuncSignature ParseArgMax; // Declaration of fused operators extern ParserFuseFuncSignature ParseFuseConvAdd; extern ParserFuseFuncSignature ParseFuseGemmRelu; @@ -357,6 +358,7 @@ RModelParser_ONNX::RModelParser_ONNX() noexcept : fOperatorsMapImpl(std::make_un RegisterOperator("ScatterND", ParseScatterND); RegisterOperator("NonZero", ParseNonZero); RegisterOperator("Clip", ParseClip); + RegisterOperator("ArgMax", ParseArgMax); } // Destructor of the parser