saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix symbolic shape inference regression in RoBERTa training (#8364)

Browse source 
* Needs to assign shape field for scalar output

* Add op test for SoftmaxCrossEntropyLoss
This commit is contained in:
KeDengMS 2021-07-13 08:15:53 -07:00 committed by GitHub
parent 7db4fc8c2a
commit eda1411e03
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 27 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
onnxruntime/python/tools/symbolic_shape_infer.py
View file

@ -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)

26
onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py
View file

@ -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):