diff --git a/onnxOpImporters.cpp b/onnxOpImporters.cpp index 521be92e..b3520a8f 100644 --- a/onnxOpImporters.cpp +++ b/onnxOpImporters.cpp @@ -5011,16 +5011,61 @@ DEFINE_BUILTIN_OP_IMPORTER(ReduceLogSumExp) RETURN_IDENTITY(inputs.at(0), node, nodeIdx); } - std::vector expResult - = unaryHelper(ctx, node, nodeIdx, inputs.at(0), nvinfer1::UnaryOperation::kEXP); + TensorOrWeights input = inputs.at(0); + TensorOrWeights inputAxes = inputs.size() >= 2 ? inputs.at(1) : TensorOrWeights(); - // Include the axes input if present to ensure reduction is performed correctly - if (inputs.size() >= 2) + OnnxAttrs attrs(node, ctx); + bool const keepdims = attrs.get("keepdims", 1); + int32_t const ndim = input.shape().nbDims; + + std::vector axes; + if (attrs.count("axes")) + { + axes = attrs.get>("axes"); + } + else if (!inputAxes.isNullTensor()) + { + ONNXTRT_CHECK_NODE( + inputAxes.is_weights(), "Axis input must be an initializer!", node, nodeIdx, ErrorCode::kUNSUPPORTED_NODE); + weightsToVector(inputAxes.weights(), &axes); + } + if (axes.empty()) + { + if (attrs.get("noop_with_empty_axes", 0) == 1) + { + RETURN_IDENTITY(inputs.at(0), node, nodeIdx); + } + axes.resize(ndim); + std::iota(axes.begin(), axes.end(), 0); + } + for (int32_t& axis : axes) { - expResult.push_back(inputs.at(1)); + convertAxis(axis, ndim, node, nodeIdx); } + std::sort(axes.begin(), axes.end()); - return importReduceLogSum(ctx, node, nodeIdx, expResult); + auto maxResult = reduceTensor(ctx, node, nodeIdx, input, nvinfer1::ReduceOperation::kMAX, inputAxes); + nvinfer1::ITensor* reducedMax = &convertToTensor(maxResult.at(0), ctx); + + nvinfer1::ITensor* maxForSub = reducedMax; + if (!keepdims) + { + maxForSub = unsqueezeTensor(ctx, *reducedMax, axes); + } + + auto& inputTensor = convertToTensor(input, ctx); + auto* shifted = getElementWiseResult(ctx, inputTensor, *maxForSub, nvinfer1::ElementWiseOperation::kSUB); + auto* shiftedExp = getUnaryResult(ctx, *shifted, nvinfer1::UnaryOperation::kEXP); + + std::vector sumInputs{TensorOrWeights{shiftedExp}}; + if (inputs.size() >= 2) + { + sumInputs.push_back(inputAxes); + } + auto sumResult = importReduceSum(ctx, node, nodeIdx, sumInputs); + auto* logSum = getUnaryResult(ctx, convertToTensor(sumResult.at(0), ctx), nvinfer1::UnaryOperation::kLOG); + auto* stableResult = getElementWiseResult(ctx, *logSum, *reducedMax, nvinfer1::ElementWiseOperation::kSUM); + return {{stableResult}}; } DECLARE_BUILTIN_OP_IMPORTER(ReduceSumSquare); DEFINE_BUILTIN_OP_IMPORTER(ReduceL2) diff --git a/onnx_backend_test.py b/onnx_backend_test.py index f62004ad..d501e438 100644 --- a/onnx_backend_test.py +++ b/onnx_backend_test.py @@ -9,6 +9,9 @@ import unittest import onnx.backend.test +import numpy as np +from onnx import TensorProto +from onnx import helper as onnx_helper import onnx_tensorrt.backend as trt @@ -176,6 +179,69 @@ # dilations not supported in ConvTRanspose layer backend_test.exclude(r'.*test_convtranspose_dilations_custom_cuda') + +class TensorRTCustomReduceLogSumExpTest(unittest.TestCase): + def test_reduce_log_sum_exp_large_finite_input_custom(self): + node = onnx_helper.make_node("ReduceLogSumExp", ["x"], ["y"], keepdims=1) + graph = onnx_helper.make_graph( + [node], + "reduce_log_sum_exp_large_finite_input_custom", + [onnx_helper.make_tensor_value_info("x", TensorProto.FLOAT, [4])], + [onnx_helper.make_tensor_value_info("y", TensorProto.FLOAT, [1])], + ) + model = onnx_helper.make_model(graph, opset_imports=[onnx_helper.make_opsetid("", 18)]) + model.ir_version = 10 + + x = np.array([250.0, 248.0, 255.0, 251.0], dtype=np.float32) + expected = np.log(np.sum(np.exp(x - np.max(x)))) + np.max(x) + + outputs = trt.run_model(model, [x], device="CUDA:0") + actual = float(outputs.y[0]) + + self.assertTrue(np.isfinite(actual)) + self.assertAlmostEqual(actual, float(expected), places=5) + + def test_reduce_log_sum_exp_keepdims_zero_custom(self): + node = onnx_helper.make_node("ReduceLogSumExp", ["x", "axes"], ["y"], keepdims=0) + graph = onnx_helper.make_graph( + [node], + "reduce_log_sum_exp_keepdims_zero_custom", + [onnx_helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 4])], + [onnx_helper.make_tensor_value_info("y", TensorProto.FLOAT, [2])], + initializer=[onnx_helper.make_tensor("axes", TensorProto.INT64, [1], np.array([1], dtype=np.int64))], + ) + model = onnx_helper.make_model(graph, opset_imports=[onnx_helper.make_opsetid("", 18)]) + model.ir_version = 10 + + x = np.array([[250.0, 248.0, 255.0, 251.0], [100.0, 101.0, 99.0, 97.0]], dtype=np.float32) + expected = np.log(np.sum(np.exp(x - np.max(x, axis=1, keepdims=True)), axis=1)) + np.max(x, axis=1) + + outputs = trt.run_model(model, [x], device="CUDA:0") + actual = np.asarray(outputs.y) + + self.assertTrue(np.all(np.isfinite(actual))) + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + + def test_reduce_log_sum_exp_noop_empty_axes_custom(self): + node = onnx_helper.make_node("ReduceLogSumExp", ["x", "axes"], ["y"], keepdims=0, noop_with_empty_axes=1) + graph = onnx_helper.make_graph( + [node], + "reduce_log_sum_exp_noop_empty_axes_custom", + [onnx_helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 2])], + [onnx_helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 2])], + initializer=[onnx_helper.make_tensor("axes", TensorProto.INT64, [0], np.array([], dtype=np.int64))], + ) + model = onnx_helper.make_model(graph, opset_imports=[onnx_helper.make_opsetid("", 18)]) + model.ir_version = 10 + + x = np.array([[250.0, 248.0], [255.0, 251.0]], dtype=np.float32) + + outputs = trt.run_model(model, [x], device="CUDA:0") + actual = np.asarray(outputs.y) + + self.assertEqual(actual.size, x.size) + np.testing.assert_allclose(actual.reshape(x.shape), x, rtol=1e-6, atol=1e-6) + globals().update(backend_test .enable_report() .test_cases)