From eda1411e034731fb35ecc8c6d9749e573dbfc1be Mon Sep 17 00:00:00 2001 From: KeDengMS Date: Tue, 13 Jul 2021 08:15:53 -0700 Subject: [PATCH] Fix symbolic shape inference regression in RoBERTa training (#8364) * Needs to assign shape field for scalar output * Add op test for SoftmaxCrossEntropyLoss --- .../python/tools/symbolic_shape_infer.py | 1 + ...untime_test_python_symbolic_shape_infer.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 869a048511..2c50b4a6b0 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -1331,6 +1331,7 @@ class SymbolicShapeInference: vi = self.known_vi_[node.output[0]] elem_type = self.known_vi_[node.input[0]].type.tensor_type.elem_type vi.type.tensor_type.elem_type = elem_type + vi.type.tensor_type.shape.CopyFrom(onnx.TensorShapeProto()) if len(node.output) > 1: data_shape = self._get_shape(node, 0) diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py index 6ba3a091ac..dd4ed8b437 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py +++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py @@ -134,6 +134,32 @@ class TestSymbolicShapeInferenceForOperators(unittest.TestCase): ] self._check_shapes(graph, inferred.graph, expected_shapes) + def test_softmax_cross_entropy_loss(self): + hidden_size = 1024 + + nodes = [ + helper.make_node("SoftmaxCrossEntropyLoss", + inputs=["logits", "labels"], + outputs=["loss"]), + ] + + inputs = [ + helper.make_tensor_value_info('logits', TensorProto.FLOAT, ['b', 's', hidden_size]), + helper.make_tensor_value_info('labels', TensorProto.INT32, ['b', 's']), + ] + + outputs = [ + helper.make_tensor_value_info('loss', TensorProto.FLOAT, None), + ] + + graph = helper.make_graph(nodes, "SoftmaxCrossEntropyLoss_Test", inputs, outputs, []) + model = helper.make_model(graph) + + inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True) + expected_shapes = [ + helper.make_tensor_value_info('loss', TensorProto.FLOAT, []) + ] + self._check_shapes(graph, inferred.graph, expected_shapes) class TestSymbolicShapeInferenceForSlice(unittest.TestCase): def check_slice_of_concat(self, input_dims, start, end, step, expected_output_dim):