def replace_op(self, graph: Graph, node: Node):
input_node = node.in_nodes()[0]
memory_pair_id = unique_id('id')
# Memory(in)
input_memory = Memory(
graph, {
'name':
'prev_splice_memory',
'id':
memory_pair_id,
'index':
1,
'size':
2,
'shape':
np.array(([input_node.shape[1] * len(node.context)]),
dtype=np.int64)
}).create_node()
# Memory(in) \
# Crop
# Input(temp) /
crop = Crop(
graph, {
'name':
'Splice_Crop',
'axis':
np.array([1], dtype=np.int64),
'offset':
np.array([input_node.shape[1]], dtype=np.int64),
'dim':
np.array([input_node.shape[1] * (len(node.context) - 1)],
dtype=np.int64)
}).create_node([input_memory])
# Crop \
# Concat
# Input /
concat_node = Concat(graph, {
'name': 'Splice_Concat',
'in_ports_count': 2,
'axis': 1
}).create_node([crop, input_node])
# Concat -> Memory(out)
Memory(
graph, {
'name':
'out_splice_memory',
'id':
memory_pair_id,
'index':
0,
'size':
2,
'shape':
np.array([input_node.shape[1] * len(node.context)],
dtype=np.int64)
}).create_node([concat_node])
return [concat_node.id]
def replace_pattern(graph: Graph, match: dict):
node = match['op']
in_shape = node.in_port(0).data.get_shape().copy()
memory_element = in_shape[1] - node.const_dim
memory_size = memory_element * len(node.context)
memory_pair_id = unique_id('id')
# Memory(in)
input_memory = Memory(
graph, {
'name': 'prev_splice_memory',
'id': memory_pair_id,
'index': 1,
'size': 2,
'shape': int64_array([memory_size])
}).create_node()
# Memory(in) \
# Crop
# Input(temp) /
crop = Crop(
graph, {
'name': 'Splice_Crop',
'axis': int64_array([1]),
'offset': int64_array([memory_element]),
'dim': int64_array([memory_size - memory_element])
}).create_node()
crop.in_port(0).connect(input_memory.out_port(0))
# Crop \
# Concat
# Input /
concat_node = Concat(graph, {
'name': 'Splice_Concat',
'in_ports_count': 2,
'axis': 1
}).create_node()
concat_node.in_port(0).connect(crop.out_port(0))
# Concat -> Memory(out)
mem_out = Memory(
graph, {
'name': 'out_splice_memory',
'id': memory_pair_id,
'index': 0,
'size': 2,
'shape': int64_array([memory_size])
}).create_node()
mem_out.in_port(0).connect(concat_node.out_port(0))
Result(graph).create_node().in_port(0).connect(mem_out.out_port(0))
if node.const_dim != 0:
memory_element_constdim = node.const_dim
memory_size_constdim = memory_element_constdim * len(node.context)
split = create_op_with_const_inputs(
graph, VariadicSplit, {
1: int64_array(1),
2: int64_array([memory_element, memory_element_constdim])
}, {
'name': node.id + '_split_const',
'out_ports_count': 2
})
split.out_port(0).connect(concat_node.in_port(1))
# create separate splice construction for const_dim
memory_pair_id = unique_id('memory_for_const_dim')
input_memory_const_dim = Memory(
graph, {
'name': 'const_dim_in_memory',
'id': memory_pair_id,
'index': 1,
'size': 2,
'shape': int64_array([memory_size_constdim])
}).create_node()
crop_const_dim = Crop(
graph, {
'name':
'const_dim_crop',
'axis':
int64_array([1]),
'offset':
int64_array([memory_element_constdim]),
'dim':
int64_array(
[memory_size_constdim - memory_element_constdim])
}).create_node()
crop_const_dim.in_port(0).connect(
input_memory_const_dim.out_port(0))
concat_node_const_dim = Concat(graph, {
'name': 'const_dim_concat',
'in_ports_count': 2,
'axis': 1
}).create_node()
concat_node_const_dim.in_port(0).connect(
crop_const_dim.out_port(0))
mem_out_const_dim = Memory(
graph, {
'name': 'const_dim_out_memory',
'id': memory_pair_id,
'index': 0,
'size': 2,
'shape': int64_array([memory_size_constdim])
}).create_node()
mem_out_const_dim.in_port(0).connect(
concat_node_const_dim.out_port(0))
Result(graph).create_node().in_port(0).connect(
mem_out_const_dim.out_port(0))
# connect splice to Split as begin and Concat as the end
split.out_port(1).connect(concat_node_const_dim.in_port(1))
crop_first = Crop(
graph, {
'name': 'const_dim_crop_first',
'axis': int64_array([1]),
'offset': int64_array([0]),
'dim': int64_array([memory_element_constdim])
}).create_node()
crop_first.in_port(0).connect(concat_node_const_dim.out_port(0))
concat_const = Concat(graph, {
'name': node.id + '_concat_const',
'axis': 1,
'in_ports_count': 2
}).create_node()
concat_const.in_port(1).connect(crop_first.out_port(0))
concat_const.in_port(0).connect(concat_node.out_port(0))
node.in_port(0).get_connection().set_destination(split.in_port(0))
node.out_port(0).get_connection().set_source(
concat_const.out_port(0))
else:
node.in_port(0).get_connection().set_destination(
concat_node.in_port(1))
node.out_port(0).get_connection().set_source(
concat_node.out_port(0))
# to avoid re-inference of shape and touching in next replacements
graph.remove_node(node.id)