def _infer_op_shapes(op_name,
attribs,
input_shapes,
output_counts,
custom_shapes={}):
func = _StandardShapeFuncs.get(op_name)
if func is None:
func = custom_shapes.get(op_name)
if func is None:
raise nnef.Error(
"shape inference function is not defined for operation '{}'".
format(op_name))
try:
output_shapes = func(*input_shapes, **attribs)
if not isinstance(output_shapes, tuple):
output_shapes = (output_shapes, )
assert len(output_counts) == len(output_shapes), \
"number of shapes ({}) does not match number of outputs ({})".format(len(output_counts), len(output_shapes))
for count, shape in zip(output_counts, output_shapes):
if isinstance(count, list):
assert isinstance(shape, list), "expected list of shapes"
assert count == len(shape), \
"number of shapes ({}) does not match number of outputs ({})".format(count, len(shape))
return output_shapes
except AssertionError as e:
raise nnef.Error(
"while inferring output shape of operation '{}': {}".format(
op_name, e))
def infer_shapes(graph, custom_shapes={}):
# type: (nnef.Graph, dict)->None
for op in graph.operations:
func = _StandardShapeFuncs.get(op.name)
if func is None:
func = custom_shapes.get(op.name)
if func is None:
raise nnef.Error("shape inference function is not defined for operation '{}'".format(op.name))
input_shapes = [_get_shape(graph, input) for input in op.inputs.values()]
try:
output_shapes = func(*input_shapes, **op.attribs)
if not isinstance(output_shapes, tuple):
output_shapes = (output_shapes,)
outputs = op.outputs.values()
assert len(outputs) == len(output_shapes), \
"number of shapes ({}) does not match number of outputs ({})".format(len(outputs), len(output_shapes))
for output, shape in zip(outputs, output_shapes):
if isinstance(output, list):
assert isinstance(shape, list), "expected list of shapes"
assert len(output) == len(shape), \
"number of shapes ({}) does not match number of outputs ({})".format(len(output), len(shape))
_set_shape(graph, output, shape)
except AssertionError as e:
raise nnef.Error("while inferring shape of tensor(s) '{}' (operation '{}'): {}"
.format(', '.join(op.outputs.values()), op.name, e))
def infer_shapes(graph, external_shapes={}, custom_shapes={}):
# type: (nnef.Graph, dict)->None
for op in graph.operations:
func = _StandardShapeFuncs.get(op.name)
if func is None:
func = custom_shapes.get(op.name)
if func is None:
raise nnef.Error(
"shape inference function is not defined for operation '{}'".
format(op.name))
if op.name == 'external':
id = op.outputs['output']
override = external_shapes.get(id)
if override is not None:
original = op.attribs['shape']
assert len(override) == len(original), \
"overridden external shape rank ({}) does not match original rank ({})".format(len(override), len(original))
_set_shape(graph, id, override)
continue
input_shapes = [
_get_shape(graph, input) for input in op.inputs.values()
]
try:
output_shapes = func(*input_shapes, **op.attribs)
if not isinstance(output_shapes, tuple):
output_shapes = (output_shapes, )
outputs = op.outputs.values()
assert len(outputs) == len(output_shapes), \
"number of shapes ({}) does not match number of outputs ({})".format(len(outputs), len(output_shapes))
for output, shape in zip(outputs, output_shapes):
if isinstance(output, list):
assert isinstance(shape, list), "expected list of shapes"
assert len(output) == len(shape), \
"number of shapes ({}) does not match number of outputs ({})".format(len(output), len(shape))
_set_shape(graph, output, shape)
except AssertionError as e:
raise nnef.Error(
"while inferring shape of tensor(s) '{}' (operation '{}'): {}".
format(', '.join(op.outputs.values()), op.name, e))
for tensor in graph.tensors.values():
if tensor.quantization:
for key, value in tensor.quantization.items():
if isinstance(value, np.ndarray):
assert _broadcastable(value.shape, tensor.shape)