def get_layer_inbound_nodes(layer_parent_pairs):
"""Get layer's inbound nodes.
The config of a layer does not include connectivity information,
nor the layer class name. These are handled by keras.Model.
So we extract them from model's config and associate them to the
corresponding layer.
"""
layer_inbound_nodes = {}
model = None
# Get a keras model which is a top-level layer.
for layer, parent_layer in layer_parent_pairs:
if parent_layer is None:
model = layer
break
if getattr(model, '_is_graph_network', None):
# Only graph network has get_config.
model_config = model.get_config()
logging.vlog(4, 'model_config: {}'.format(model_config))
if 'layers' in model_config:
layers_config = model_config['layers']
for config in layers_config:
if 'inbound_nodes' in config:
layer_inbound_nodes[
config['name']] = config['inbound_nodes']
return layer_inbound_nodes
def func_graph_to_nndct(func_graph, scope_to_layer=None):
# op_name => Node name
tf_graph, input_signature, structured_output_tensors = func_graph
computation_graph = ComputationGraph.from_tf_graph(tf_graph,
scope_to_layer)
logging.vlog(2, 'ComputationGraph\n {}'.format(computation_graph))
# Parse computation nodes to nndct nodes.
nndct_nodes = []
for node in computation_graph.nodes:
nndct_nodes.extend(converter.convert(node))
# Create all tensors
tensors = {}
for node in nndct_nodes:
for name in node.output_names:
tensors[name] = node.produce(name)
# Build connections.
for node in nndct_nodes:
for name in node.input_names:
node.consume(tensors[name])
graph = ops.Graph()
for node in nndct_nodes:
graph.add_node(node)
# The tensors in FuncGraph.structured_output_tensors are outputs from
# Identity node added to the graph. Since all Identity nodes will be removed
# in graph refining, so we have to find the actual output tensors before
# that process.
output_tensors = []
for tensor in nest.flatten(structured_output_tensors):
# The output tensors does not exist in graph, so we can't get the output
# node by tensor's producer, like:
# node = graph.tensor(tf_tensor.name).producer
node_name = tf_utils.node_name_from_input(tensor.name)
node = graph.node(node_name)
assert node.op.type == OpTypes.IDENTITY
output_tensors.append(node.in_tensors[0])
output_tensors = nest.pack_sequence_as(structured_output_tensors,
output_tensors)
utils.maybe_export_graph('{}/{}'.format(_EXPORT_DIR, _RAW_NNDCT_GRAPH),
graph)
graph = run_graph_refining(graph)
utils.maybe_export_graph('{}/{}'.format(_EXPORT_DIR, _FINAL_NNDCT_GRAPH),
graph)
logging.vlog(2, 'NndctGraph before sorting:\n{}'.format(graph))
graph = utils.topological_sort(graph)
# Get args part from input_signature (args, kwargs)
graph.input_signature = input_signature[0]
graph.structured_output_tensors = output_tensors
logging.vlog(2, 'NndctGraph\n{}'.format(graph))
logging.vlog(2, 'input_signature:{}'.format(input_signature))
logging.vlog(2, 'output_tensors:{}'.format(output_tensors))
return graph
def run_graph_refining(graph):
# Executed in sequence.
refiners = [
FoldConst,
FoldBias,
RemoveIdentity,
RemoveRNNRedundantInput,
RemoveIsolatedNode,
MergeBidirectionalRNN,
RenameParamTensor,
SetAttrForBinaryOp,
]
for refiner_cls in refiners:
refiner = refiner_cls()
result = refiner.refine_graph(graph)
logging.vlog(
2, 'Refining pass [{}]: {}'.format(result.refiner, result.message))
return graph
def run_graph_refining(graph):
# Executed in sequence.
refiners = [
FoldConstRefiner,
FoldBiasRefiner,
RemoveIdentityRefiner,
RemoveConstRefiner,
RemoveIsolatedRefiner,
MergeBidirectionalRefiner,
RenameParamTensorRefiner,
]
for refiner_cls in refiners:
refiner = refiner_cls()
result = refiner.refine_graph(graph)
logging.vlog(
2,
'Result of refining pass [{}]: {}'.format(result.refiner,
result.message))
return graph
def from_keras_model(model, input_signature=None):
"""Trace model call to get a func graph and convert that func graph
to nndct graph.
"""
logging.vlog(1, 'input_signature: {}'.format(input_signature))
if not generic_utils.is_list_or_tuple(input_signature):
input_signature = generic_utils.to_list(input_signature)
func_graph = get_func_graph(model, input_signature)
scope_to_layer = map_scope_to_layer(model)
logging.vlog(
1, 'scope_name: (layer, parent_layer)\n{}'.format('\n'.join(
[f'{key}: {value}' for key, value in scope_to_layer.items()])))
graph = func_graph_to_nndct(func_graph, scope_to_layer)
graph.name = model.name
graph.data_format = keras_utils.data_format()
return graph
def from_keras_model(model, input_signature):
logging.vlog(1, 'input_signature: {}'.format(input_signature))
#input_signature = [tf.TensorSpec(shape=batch_input_shape, dtype=dtype)]
if not generic_utils.is_list_or_tuple(input_signature):
input_signature = generic_utils.to_list(input_signature)
flat_input_signature = nest.flatten(input_signature)
batch_input_signature = []
for signature in flat_input_signature:
batch_input_signature.append(
tf.TensorSpec(shape=(1, ) + signature.shape,
dtype=signature.dtype))
batch_input_signature = nest.pack_sequence_as(input_signature,
batch_input_signature)
func_graph = get_func_graph(model, batch_input_signature)
scope_to_layer = map_scope_to_layer(model, '')
logging.vlog(1, 'scope_to_layer:\n{}'.format(scope_to_layer))
graph = parse_to_graph(func_graph, scope_to_layer)
graph.name = model.name
return graph
def parse_to_graph(func_graph, scope_to_layer=None):
# op_name => ComputationNode name
tf_graph, input_signature, structured_output_tensors = func_graph
op_to_node = {}
nodes = {}
for op in tf_graph.get_operations():
layer = None
if scope_to_layer:
layer = belongs_to_keras_layer(op, scope_to_layer)
# keras.layers or tf.Operation as a node
node = layer if layer else op
if node.name not in nodes:
nodes[node.name] = ComputationNode(node)
nodes[node.name].scope_ops.append(op)
op_to_node[op.name] = node.name
mark_computation_edges(nodes, op_to_node)
for node in nodes.values():
logging.vlog(2, 'ComputationNode\n {}'.format(str(node)))
# Parse computation nodes to nndct nodes.
nndct_nodes = []
for node in nodes.values():
nndct_nodes.extend(nest.flatten(converter.convert(node)))
# Create all tensors
tensors = {}
for node in nndct_nodes:
for name in node.output_names:
tensors[name] = node.produce(name)
# Build connections.
for node in nndct_nodes:
for name in node.input_names:
node.consume(tensors[name])
graph = ops.Graph()
for node in nndct_nodes:
graph.add_node(node)
# The tensors in FuncGraph.structured_output_tensors are outputs from
# Identity node added to the graph. Since all Identity nodes will be removed
# in graph refining, so we have to find the actual output tensors before
# that process.
output_tensors = []
for tensor in nest.flatten(structured_output_tensors):
# The output tensors does not exist in graph, so we can't get the output
# node by tensor's producer, like:
# node = graph.tensor(tf_tensor.name).producer
node_name = tf_utils.node_name_from_input(tensor.name)
node = graph.node(node_name)
assert node.op.type == OpTypes.IDENTITY
output_tensors.append(node.in_tensors[0])
output_tensors = nest.pack_sequence_as(structured_output_tensors,
output_tensors)
utils.maybe_export_graph(_RAW_NNDCT_GRAPH, graph)
graph = run_graph_refining(graph)
utils.maybe_export_graph(_FINAL_NNDCT_GRAPH, graph)
graph = utils.topological_sort(graph)
# Get args part from input_signature (args, kwargs)
graph.input_signature = input_signature[0]
graph.structured_output_tensors = output_tensors
logging.vlog(2, str(graph))
logging.vlog(2, 'input_signature:{}'.format(input_signature))
logging.vlog(2, 'output_tensors:{}'.format(output_tensors))
return graph
def build(self, filepath, quantized=False, as_layer=False):
class_name = self._graph.name
base_class = keras.Model
call_fn_name = 'call'
if as_layer:
if quantized:
base_class = base_layer.Layer
call_fn_name = '_internal_call'
else:
base_class = keras.layers.Layer
class_spec = ClassSpec(class_name, base_class, call_fn_name, quantized)
writer = writer_lib.ModuleClassWriter(self._graph, class_spec)
layer_to_node = writer.write(filepath)
# TODO(yuwang): Use code below.
#py_module_name = "_".join(["nndct", module_name])
#spec = importlib.util.spec_from_file_location(py_module_name, filepath)
#py_module = importlib.util.module_from_spec(spec)
#sys.modules[py_module_name] = py_module
#spec.loader.exec_module(py_module)
#rebuilt_module = py_module.__dict__[module_name]()
loaded_module = imp.load_source('nndct_rebuilt_model', filepath)
rebuilt_model = getattr(loaded_module, class_name)()
dummy_inputs = []
for spec in nest.flatten(self._graph.input_signature):
logging.vlog(1, spec)
dummy_inputs.append(array_ops.ones(spec.shape, dtype=spec.dtype))
dummy_inputs = nest.pack_sequence_as(self._graph.input_signature,
dummy_inputs)
#input_data = dummy_inputs if len(dummy_inputs) > 1 else dummy_inputs[0]
rebuilt_model(*dummy_inputs)
layer_nodes = []
# Reload weights
for layer_name, node in layer_to_node.items():
layer = getattr(rebuilt_model, layer_name)
# If there is a ParamName definition, then map ParamName's member to
# keras layer's param; If there is no ParamName,
# then export params in the order they are saved in the op.
weights = []
if hasattr(node.op, 'ParamName'):
params = keras_utils.keras_layer_params(layer)
for name in params.keys():
param = utils.op_param_by_name(node.op, name)
if not param:
continue
ndarray = tensor_utils.to_tf_numpy(
node.op.get_param(param))
weights.append(ndarray)
logging.vlog(
2, 'Reload weights of {}: name={}, shape={}'.format(
layer.name, name, ndarray.shape))
else:
for name, tensor in node.op.params.items():
ndarray = tensor_utils.to_tf_numpy(tensor)
weights.append(ndarray)
logging.vlog(
2, "Reload weights of {}: name={}, shape={}".format(
layer.name, name, ndarray.shape))
layer_nodes.append((layer, node))
if weights:
layer.set_weights(weights)
return rebuilt_model, layer_nodes