natas_hw
/
openvino
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
openvino/model-optimizer/mo/graph/connection.py

247 lines
9.9 KiB
Python
Raw Blame History

"""
Copyright (C) 2018-2020 Intel Corporation
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
from collections import namedtuple
from copy import deepcopy
from mo.utils.error import Error
class Connection:
def __init__(self, graph, source, destinations: list, control_flow=False):
self.graph = graph
self.source = source
self.destinations = destinations
self.control_flow = control_flow
self.data = namedtuple('Data', ['get_value', 'get_shape'])
self.data.get_value = self._get_value
self.data.get_shape = self._get_shape
def _get_value(self):
if self.graph.stage == 'front':
return None
return self.source.node.out_node(self.source.idx, control_flow=self.control_flow).value
def _get_shape(self):
if self.graph.stage == 'front':
return None
return self.source.node.out_node(self.source.idx, control_flow=self.control_flow).shape
def get_source(self):
return self.source
def get_destination(self):
if self.destinations and len(self.destinations) > 1:
raise Error("Connection has more than one destination: {}".format(len(self.destinations)))
return self.destinations[0] if self.destinations else None
def get_destinations(self):
return self.destinations
def set_source(self, port):
# In this method we are changing source for a connection with given port.
# See detailed example below.
#
# SOURCE - Op1(out_port:0)
#
# | Op4(in_port:0)
# DESTINATIONS - | Op3(in_port:0)
# | Op2(in_port:0)
#
# NEW PORT - Op5(out_port:0)
#
# ,--->Op4(in_port:0)
# CONNECTION ,--->Op3(in_port:0)
# Op1(out_port:0)--->Op2(in_port:0)
#
# When we set source for connection we disconnect existing source and reconnect all consumers to
# the new given port with type='out'.
#
# UPDATED CONNECTION ,--->Op4(in_port:0)
# ,--->Op3(in_port:0)
# Op5(out_port:0)--->Op2(in_port:0)
#
if port.type == 'in':
raise Error("Wrong port type in set_source method. Should be 'out' but given 'in'")
if self.control_flow is True:
raise Error("Cannot operate with connection with control_flow=True")
if self.graph.stage == 'front':
scr_node = port.node
# Reconnecting all destinations as consumers to the source port preserving edge attrs
for dst_port in self.destinations:
edge_attrs, u, v, key = dst_port.get_in_edge_attrs(data=True)
if u is not None:
edge_attrs['out'] = port.idx
self.graph.remove_edge(u, v, key=key)
self.graph.add_edge(scr_node.id, v, **edge_attrs)
else:
self.graph.create_edge(scr_node, dst_port.node, port.idx, dst_port.idx)
else:
# Create out data node if not exists and mark node with need_shape_inference = True
# In case if data node exists just use it.
port._create_data_if_necessary()
port_out_data = port.node.out_node(port.idx)
if self.source is not None and self.source.idx in self.source.node.out_nodes():
source_out_data = self.source.node.out_node(self.source.idx)
# Copy attrs from source_out_data to port_out_data
attrs = deepcopy(source_out_data.attrs())
# Remove debug info
if 'fw_tensor_debug_info' in source_out_data.attrs():
del self.graph.node[source_out_data.id]['fw_tensor_debug_info']
# Copy attrs to new data node
for attr in attrs:
port_out_data[attr] = attrs[attr]
for dst_port in self.destinations:
edge_attrs, u, v, key = dst_port.get_in_edge_attrs(data=True)
if u is not None:
self.graph.remove_edge(u, v, key=key)
self.graph.add_edge(port_out_data.id, v, **edge_attrs)
else:
self.graph.add_edge(port_out_data.id, dst_port.node.id, **{'in': dst_port.idx})
def set_destination(self, port):
# In this method we are changing destination for a connection with given port with type 'in'.
# This method requires exactly one destination or empty destinations list.
# See detailed example below.
#
# SOURCE - Op1(out_port:0)
#
# DESTINATIONS - Op2(in_port:0)
#
# NEW PORT - Op3(in_port:0)
#
# CONNECTION
# Op1(out_port:0)--->Op2(in_port:0)
#
# When we set destination for connection we disconnect destination port if exists and connect source to
# the new given port with type='in'.
#
# UPDATED CONNECTION
#
# Op1(out_port:0)--->Op3(in_port:0)
#
def check_and_remove_edge():
if self.destinations:
for destination in self.destinations:
edge_attrs, u, v, key = destination.get_in_edge_attrs(data=True)
if u is None:
raise Error(
"Broken Connection object! Destination (node:{}) is not connected to source.".format(
destination.node.name))
destination.disconnect()
return edge_attrs
return None
if self.destinations and len(self.destinations) > 1:
raise Error("set_destination applicable only for connections that has exactly one destination or \
when there is no destinations")
if port.type == 'out':
raise Error("Wrong port type in set_destination method. Should be 'in' but given 'out'")
if self.control_flow is True:
raise Error("Cannot operate with connection with control_flow=True")
if self.graph.stage == 'front':
if self.source is not None:
node = self.source.node
check_and_remove_edge()
self.graph.create_edge(node, port.node, out_port=self.source.idx, in_port=port.idx)
self.destinations = [port]
else:
# create out node if not exists and mark node with need_shape_inference = True
# in case if data node exists just use it as is
if self.source is not None:
data_node = self.source._create_data_if_necessary()
edge_attrs = check_and_remove_edge() or {}
edge_attrs.update({'in': port.idx})
self.graph.add_edge(data_node.id, port.node.id, **edge_attrs)
self.destinations = [port]
def add_destination(self, port):
# In this method we are adding destination port with type 'in' for a connection.
# See detailed example below.
#
# SOURCE - Op1(out_port:0)
#
# DESTINATIONS - Op2(in_port:0)
#
# NEW PORT - Op3(in_port:0)
#
# CONNECTION
# Op1(out_port:0)--->Op2(in_port:0)
#
# When we set destination for connection we disconnect destination port if exists and connect source to
# the new given port with type='in'.
#
# UPDATED CONNECTION
# ,-->Op3(in_port:0)
# Op1(out_port:0)--->Op2(in_port:0)
#
if self.control_flow is True:
raise Error("Cannot operate with connection with control_flow=True")
if self.source is None:
raise Error("Can not add destination for connection without source port!")
if self.graph.stage == 'front':
node = self.source.node
self.graph.create_edge(node, port.node, out_port=self.source.idx, in_port=port.idx)
else:
data_node = self.source._create_data_if_necessary()
self.graph.add_edge(data_node.id, port.node.id, **{'in': port.idx})
self.destinations.append(port)
def remove(self):
# This method deletes all edges in connection. After that connection is not more accessible.
# See detailed example below.
#
# SOURCE - Op1(out_port:0)
#
# | Op4(in_port:0)
# DESTINATIONS - | Op3(in_port:0)
# | Op2(in_port:0)
#
# ,--->Op4(in_port:0)
# CONNECTION ,--->Op3(in_port:0)
# Op1(out_port:0)--->Op2(in_port:0)
#
# After removing edges connection will be empty
#
# REMOVED CONNECTION
# Op5(out_port:0) Op4(in_port:0) Op2(in_port:0) Op3(in_port:0)
#
if self.destinations:
for dst_port in self.destinations:
dst_port.disconnect()
self.source = None
self.destinations = []
def insert_node(self, new_node):
assert len(new_node.out_ports()) == 1, 'The node {} has several output ports'.format(new_node.soft_get('name'))
source_port = self.get_source()
self.set_source(new_node.out_port(0))
source_port.connect(new_node.in_port(0))
Reference in New Issue View Git Blame Copy Permalink