ONNX import: simplify logic for determining number of neurons in matmul op
@@ -802,35 +802,15 @@ def onnx_extract_matmul(graph, node, layers):
act = 'identity'
neurons = None
- # determine dimensions of the multiply
- a_shape = None
- b_shape = None
# check initializer for weight tensors
for init in graph.initializer:
- if init.name == node.input[0]:
- a_shape = numpy_helper.to_array(init).shape
if init.name == node.input[1]:
- b_shape = numpy_helper.to_array(init).shape
+ neurons = numpy_helper.to_array(init).shape[1]
- # check inferred shapes in graph
- for v in graph.value_info:
- if v.name == node.input[0]:
- a_shape = (v.type.tensor_type.shape.dim[0].dim_value,
- v.type.tensor_type.shape.dim[1].dim_value)
- if v.name == node.input[1]:
- b_shape = (v.type.tensor_type.shape.dim[0].dim_value,
- v.type.tensor_type.shape.dim[1].dim_value)
-
- if a_shape is None or b_shape is None:
+ if neurons is None:
raise OnnxParseError('Unable to determine number of neurons '
'in FC layer.')
- # set number of neurons according to shape
- if a_shape[0] == 1:
- neurons = b_shape[1]
- else:
- neurons = a_shape[0]
-
# check if bias is added by the next op
out = onnx_get_out_nodes(graph, node)
for n in out:
