diff --git a/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx b/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx index 025d6d678088a..7597ac1846cfe 100644 --- a/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx +++ b/tmva/sofie/inc/TMVA/ROperator_Softmax.hxx @@ -104,12 +104,21 @@ public: out << SP << SP << SP << "sum += y_ptr[j];\n"; out << SP << SP << "}\n"; - out << SP << SP << fType << " inv_sum = 1.0f / sum;\n"; - out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n"; - out << SP << SP << SP << "y_ptr[j] *= inv_sum;\n"; - if (fLogSoftmax) - out << SP << SP << SP << "y_ptr[j] = " << logFunction << "(y_ptr[j]);\n"; - out << SP << SP << "}\n"; + if (fLogSoftmax) { + // evaluate (x - vmax) - log(sum(exp(x - vmax))) instead of + // log(exp(x - vmax) / sum): for inputs far apart from the maximum + // exp(x - vmax) underflows to 0 and its log is -inf, while the + // correct result is a large negative number (see #23547) + out << SP << SP << fType << " log_sum = " << logFunction << "(sum);\n"; + out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n"; + out << SP << SP << SP << "y_ptr[j] = (x_ptr[j] - vmax) - log_sum;\n"; + out << SP << SP << "}\n"; + } else { + out << SP << SP << fType << " inv_sum = 1.0f / sum;\n"; + out << SP << SP << "for (int j = 0; j < " << axis_size << "; ++j) {\n"; + out << SP << SP << SP << "y_ptr[j] *= inv_sum;\n"; + out << SP << SP << "}\n"; + } out << SP << "}\n"; } else { @@ -166,20 +175,44 @@ public: for (size_t j = 0; j < size-1; j++) out << SP; out << "}\n"; // normalize - for (size_t j = 0; j < size-1; j++) out << SP; - out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n"; - for (size_t j = 0; j < size; j++) out << SP; - out << "size_t id = index + i"; - if (stride[axis].GetVal() != "1") out << "*(" << stride[axis] << ")"; - out << ";\n"; - for (size_t j = 0; j < size; j++) out << SP; - out << "tensor_" << fNY << "[id] /= sum;\n"; if (fLogSoftmax) { - for (size_t j = 0; j < size; j++) out << SP; - out << "tensor_" << fNY << "[id] = " << logFunction << "(tensor_" << fNY << "[id]);\n"; + // same numerically stable form as in the last-axis case above: + // (x - vmax) - log(sum(exp(x - vmax))), see #23547 + for (size_t j = 0; j < size - 1; j++) + out << SP; + out << fType << " log_sum = " << logFunction << "(sum);\n"; + for (size_t j = 0; j < size - 1; j++) + out << SP; + out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n"; + for (size_t j = 0; j < size; j++) + out << SP; + out << "size_t id = index + i"; + if (stride[axis].GetVal() != "1") + out << "*(" << stride[axis] << ")"; + out << ";\n"; + for (size_t j = 0; j < size; j++) + out << SP; + out << "tensor_" << fNY << "[id] = (tensor_" << fNX << "[id] - vmax) - log_sum;\n"; + for (size_t j = 0; j < size - 1; j++) + out << SP; + out << "}\n"; + } else { + for (size_t j = 0; j < size - 1; j++) + out << SP; + out << "for (int i = 0; i < " << fShape[axis] << "; i++) {\n"; + for (size_t j = 0; j < size; j++) + out << SP; + out << "size_t id = index + i"; + if (stride[axis].GetVal() != "1") + out << "*(" << stride[axis] << ")"; + out << ";\n"; + for (size_t j = 0; j < size; j++) + out << SP; + out << "tensor_" << fNY << "[id] /= sum;\n"; + for (size_t j = 0; j < size - 1; j++) + out << SP; + out << "}\n"; } - for (size_t j = 0; j < size-1; j++) out << SP; - out << "}\n"; //end loops for (int i = static_cast(k) - 1; i >= 0; i--) { for (int j = 0; j < i; j++) out << SP; diff --git a/tmva/sofie/test/TestCustomModelsFromONNX.cxx b/tmva/sofie/test/TestCustomModelsFromONNX.cxx index ad1f7286d4555..9ad1ff693a838 100644 --- a/tmva/sofie/test/TestCustomModelsFromONNX.cxx +++ b/tmva/sofie/test/TestCustomModelsFromONNX.cxx @@ -1024,6 +1024,56 @@ TEST(ONNX, Softmax4d) expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); } +TEST(ONNX, LogSoftmax1d) +{ + SofieReference ref = readReference("LogSoftmax1d"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmax1d", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +TEST(ONNX, LogSoftmax2d) +{ + SofieReference ref = readReference("LogSoftmax2d"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmax2d", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +TEST(ONNX, LogSoftmax3d) +{ + SofieReference ref = readReference("LogSoftmax3d"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmax3d", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +// Inputs more than ~88 apart from the maximum make exp(x - max) underflow to 0 +// in float32, so log(softmax(x)) returns -inf where the stable form +// (x - max) - log(sum(exp(x - max))) returns a large negative number +TEST(ONNX, LogSoftmaxLargeRange) +{ + SofieReference ref = readReference("LogSoftmaxLargeRange"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmaxLargeRange", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + +// same as above, but reducing along the first axis, which takes the generic +// code path of the SOFIE Softmax operator +TEST(ONNX, LogSoftmaxLargeRangeAxis0) +{ + SofieReference ref = readReference("LogSoftmaxLargeRangeAxis0"); + + ASSERT_INCLUDE_AND_RUN(std::vector, "LogSoftmaxLargeRangeAxis0", ref.f32("input0")); + + expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE); +} + TEST(ONNX, ConvTranspose1d) { SofieReference ref = readReference("ConvTranspose1d"); diff --git a/tmva/sofie/test/generate_input_models.py b/tmva/sofie/test/generate_input_models.py index 28a06d4ae68eb..7a84fe589997d 100644 --- a/tmva/sofie/test/generate_input_models.py +++ b/tmva/sofie/test/generate_input_models.py @@ -3532,6 +3532,103 @@ def make_Log(): return _model(graph, opset=14, ir_version=7, producer_name='pytorch', producer_version='1.13.1') +def make_LogSoftmax1d(): + """Ops: LogSoftmax""" + nodes = [ + helper.make_node("LogSoftmax", ["X"], ["Y"]), + ] + graph = helper.make_graph( + nodes, + "LogSoftmax", + inputs=[ + _vi("X", FLOAT, [3]), + ], + outputs=[ + _vi("Y", FLOAT, [3]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + +def make_LogSoftmax2d(): + """Ops: LogSoftmax""" + nodes = [ + helper.make_node("LogSoftmax", ["X"], ["Y"], axis=-1), + ] + graph = helper.make_graph( + nodes, + "LogSoftmax", + inputs=[ + _vi("X", FLOAT, [2, 3]), + ], + outputs=[ + _vi("Y", FLOAT, [2, 3]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + +def make_LogSoftmax3d(): + """Ops: LogSoftmax""" + nodes = [ + helper.make_node("LogSoftmax", ["X"], ["Y"], axis=1), + ] + graph = helper.make_graph( + nodes, + "LogSoftmax", + inputs=[ + _vi("X", FLOAT, [2, 3, 4]), + ], + outputs=[ + _vi("Y", FLOAT, [2, 3, 4]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + +def make_LogSoftmaxLargeRange(): + """Ops: LogSoftmax + + The inputs of each row are 100 apart, so exp(x - max) underflows to 0 in + float32 for all but the largest one. log(softmax(x)) is -inf there, while + the stable form (x - max) - log(sum(exp(x - max))) stays finite.""" + nodes = [ + helper.make_node("LogSoftmax", ["X"], ["Y"], axis=-1), + ] + graph = helper.make_graph( + nodes, + "LogSoftmax", + inputs=[ + _vi("X", FLOAT, [2, 3]), + ], + outputs=[ + _vi("Y", FLOAT, [2, 3]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + +def make_LogSoftmaxLargeRangeAxis0(): + """Ops: LogSoftmax + + Same as LogSoftmaxLargeRange but with the reduction along the first axis, + which takes the generic code path of the SOFIE Softmax operator.""" + nodes = [ + helper.make_node("LogSoftmax", ["X"], ["Y"], axis=0), + ] + graph = helper.make_graph( + nodes, + "LogSoftmax", + inputs=[ + _vi("X", FLOAT, [3, 2]), + ], + outputs=[ + _vi("Y", FLOAT, [3, 2]), + ], + ) + return _model(graph, opset=17, ir_version=8) + + def make_MatMul_1D_Constant(): """Ops: MatMul""" nodes = [ @@ -5955,6 +6052,11 @@ def make_IdentityWeightOutput(): "Linear_32": make_Linear_32, "Linear_64": make_Linear_64, "Log": make_Log, + "LogSoftmax1d": make_LogSoftmax1d, + "LogSoftmax2d": make_LogSoftmax2d, + "LogSoftmax3d": make_LogSoftmax3d, + "LogSoftmaxLargeRange": make_LogSoftmaxLargeRange, + "LogSoftmaxLargeRangeAxis0": make_LogSoftmaxLargeRangeAxis0, "MatMul_1D_Constant": make_MatMul_1D_Constant, "MatMul_Stacked": make_MatMul_Stacked, "MatMul_Stacked2": make_MatMul_Stacked2, @@ -6862,6 +6964,13 @@ def rand_f32(seed, shape): "Linear_32": [f32(np.full(3200, 1.0), (32, 100))], "Linear_64": [f32(np.full(6400, 1.0), (64, 100))], "Log": [f32([1.0, 2.0, 3.0, 4.0], (4,))], + "LogSoftmax1d": [f32([-1.0, 0.0, 1.0], (3,))], + "LogSoftmax2d": [f32([-1.0, 0.0, 1.0, 2.0, 1.0, 0.0], (2, 3))], + "LogSoftmax3d": [rand_f32(11, (2, 3, 4))], + # 100 apart per row: exp(x - max) underflows to 0 in float32 away from the + # maximum, which log(softmax(x)) turns into -inf + "LogSoftmaxLargeRange": [f32([0.0, 100.0, 200.0, 200.0, 100.0, 0.0], (2, 3))], + "LogSoftmaxLargeRangeAxis0": [f32([0.0, 100.0, 100.0, 200.0, 200.0, 300.0], (3, 2))], "Max": [ f32([1.0, 2.0, -1.0], (1, 3)), f32([3.0, 0.0, 4.0], (1, 3)), @@ -7698,6 +7807,23 @@ def resolve(name): return [y.astype(np.float32)] +def _logsoftmax_reference(model, feeds): + """LogSoftmax evaluated with the numerically stable formula in double + precision. + + The ReferenceEvaluator computes log(softmax(x)), whose exp(x - max) + underflows to 0 in float32 once the inputs are more than ~88 apart from the + maximum, so it yields -inf where the correct result is a large negative + number (see root-project/root#23547).""" + node = next(n for n in model.graph.node if n.op_type == "LogSoftmax") + axis = next((a.i for a in node.attribute if a.name == "axis"), -1) + x = np.asarray(feeds[node.input[0]], dtype=np.float64) + axis = axis % x.ndim + shifted = x - x.max(axis=axis, keepdims=True) + y = shifted - np.log(np.exp(shifted).sum(axis=axis, keepdims=True)) + return [y.astype(np.float32)] + + EXPECTED_OVERRIDES = { "GRUBidirectional": _recurrent_reference, "LSTMBidirectional": _recurrent_reference, @@ -7712,6 +7838,8 @@ def resolve(name): "IdentityWeightBatchNorm": _batchnorm_reference, "BatchNormEpsilon": _batchnorm_reference, "BatchNormReluEpsilon": _batchnorm_reference, + "LogSoftmaxLargeRange": _logsoftmax_reference, + "LogSoftmaxLargeRangeAxis0": _logsoftmax_reference, }