def infer_shapes( graph, # type: ONNXGraph source_shapes=None, # type: typing.Union[typing.Dict[str, typing.List[int]], typing.List[int], int, None] custom_shapes=None, # type: typing.Optional[typing.Dict[str, typing.Callable]] ): # type: (...)->None shape_functions = dict(_DefaultShapes) if custom_shapes: shape_functions.update(custom_shapes) graph.sort() shape_fixer.fix_input_shapes(graph, source_shapes) for op in graph.operations: # Shape prop assert op.name in shape_functions, "No shape function for {}".format( op.name) inferred_shapes, inferred_dtypes = shape_functions[op.name](op) assert not utils.has_le_0(inferred_shapes) assert len(inferred_shapes) == len(inferred_dtypes) == len(op.outputs) for new_shape, new_dtype, tensor in zip(inferred_shapes, inferred_dtypes, op.outputs): assert utils.compatible_shapes(tensor.shape, new_shape) tensor.shape = new_shape assert tensor.dtype is None or tensor.dtype == new_dtype tensor.dtype = new_dtype graph_utils.remove_unreachable(graph)
def evaluate_and_convert(tf_graph, source_shapes=None): # type: (TFGraph, typing.Union[typing.Dict[str, typing.List[int]], typing.List[int], int, None])->None tf_graph.sort() if isinstance(source_shapes, dict): source_shapes = {(k + ':0' if ':' not in k else k): v for k, v in six.iteritems(source_shapes)} shape_fixer.fix_input_shapes(tf_graph, source_shapes) const_value_by_tensor = {} for tensor in tf_graph.tensors: if tensor.is_constant: const_value_by_tensor[tensor] = tf_pb_eval._evaluate_constant( tensor) elif tensor.is_variable: const_value_by_tensor[tensor] = tensor.data for op in tf_graph.operations: # Shape prop if op.name not in tf_pb_shape_inference._DefaultPropagators: raise utils.NNEFToolsException( "Operation '{}' is not supported".format(op.name)) propagated_shapes, propagated_dtypes = \ tf_pb_shape_inference._DefaultPropagators[op.name](op, const_value_by_tensor) assert not utils.has_le_0(propagated_shapes) assert len(propagated_shapes) == len(propagated_dtypes) == len( op.outputs) for new_shape, new_dtype, tensor in zip(propagated_shapes, propagated_dtypes, op.outputs): assert utils.compatible_shapes(tensor.shape, new_shape) tensor.shape = new_shape assert tensor.dtype is None or tensor.dtype == new_dtype tensor.dtype = new_dtype # Evaluation if op.name in tf_pb_eval._DefaultOpEvaluators: tf_pb_eval._DefaultOpEvaluators[op.name](op, const_value_by_tensor) # Conversion assert op.name in DefaultConverters, "No tf_pb_to_tf_py converter for {}".format( op.name) DefaultConverters[op.name](op, const_value_by_tensor) for tensor in tf_graph.tensors: tensor.dtype = _tf_py_dtype_by_tf_pb_dtype.get(tensor.dtype, None) for tensor in tf_graph.tensors: if tensor.is_variable: label = tensor.name if label is not None: if label.endswith(':0'): label = label[:-2] label = label.replace(':', '_') tensor.label = label
def propagate(graph, source_shapes=None): # type: (ONNXGraph, typing.Union[typing.Dict[str, typing.List[int]], typing.List[int], int, None])->None graph.sort() shape_fixer.fix_input_shapes(graph, source_shapes) for op in graph.operations: # Shape prop assert op.name in _DefaultPropagators, "No shape propagator for {}".format( op.name) propagated_shapes, propagated_dtypes = _DefaultPropagators[op.name](op) assert not utils.has_le_0(propagated_shapes) assert len(propagated_shapes) == len(propagated_dtypes) == len( op.outputs) for new_shape, new_dtype, tensor in zip(propagated_shapes, propagated_dtypes, op.outputs): assert utils.compatible_shapes(tensor.shape, new_shape) tensor.shape = new_shape assert tensor.dtype is None or tensor.dtype == new_dtype tensor.dtype = new_dtype graph_utils.remove_unreachable(graph)