def check_init_states(graph: Graph, match: dict):
"""
Check if cell have initial states and create zeros states if not.
"""
rnn_layer = match['rnn_layer']
num_directions = 2 if rnn_layer.direction == 'bidirectional' else 1
batch_size = rnn_layer.in_node(0).shape[rnn_layer.batch_dim]
h_init_port = 5
c_init_port = 6
if h_init_port not in rnn_layer.in_nodes():
h_shape = [num_directions, batch_size, rnn_layer.hidden_size] # from ONNX spec
h_init = np.full(h_shape, 0, dtype=np.float32)
Op.create_and_connect_input_data_node(
graph,
rnn_layer,
{'value': h_init, 'shape': int64_array(h_init.shape)},
{'in': h_init_port, 'permutation': None}
)
if rnn_layer.op == 'LSTM':
if c_init_port not in rnn_layer.in_nodes():
c_shape = [num_directions, batch_size, rnn_layer.hidden_size] # from ONNX spec
c_init = np.full(c_shape, 0, dtype=np.float32)
Op.create_and_connect_input_data_node(
graph,
rnn_layer,
{'value': c_init, 'shape': int64_array(c_init.shape)},
{'in': c_init_port, 'permutation': None}
)
def repack_weights(self, graph: Graph, match: dict):
# Concat W, R in IE- format
# Delete useless num_dir dimensions and n_cells dimensions in W, R, B (peepholes?)
lstm = match['rnn_layer']
W, R, B = match['W'].value.copy(), match['R'].value.copy(), match['B'].value.copy()
graph.remove_edge(match['W'].id, lstm.id)
graph.remove_edge(match['R'].id, lstm.id)
graph.remove_edge(match['B'].id, lstm.id)
# Sum component of B that correspond to W and R
if lstm.op == 'GRU' and lstm.linear_before_reset:
B_shape = np.array(B.shape)
B_shape[3] = 4
B_shape[2] = 1
B_tmp = np.zeros(shape=B_shape)
B_tmp[:, :, :, 0, :] = B[:, :, 0, 0, :] + B[:, :, 1, 0, :]
B_tmp[:, :, :, 1, :] = B[:, :, 0, 1, :] + B[:, :, 1, 1, :]
B_tmp[:, :, :, 2, :] = B[:, :, 0, 2, :][:, :, np.newaxis, :]
B_tmp[:, :, :, 3, :] = B[:, :, 1, 2, :][:, :, np.newaxis, :]
B = B_tmp
else:
B = np.sum(B, axis=2, keepdims=True)
# Concatenate W, R to IE-compatible format
assert len(W.shape) == 5
assert len(R.shape) == 5
WR = np.concatenate([W, R], axis=4)
# Squeeze useless dimensions
assert WR.shape[0] == 1 # num_dir == 1
assert WR.shape[1] == 1 # num_cells == 1
assert B.shape[0] == 1
assert B.shape[1] == 1
WR = WR.squeeze(axis=(0, 1))
B = B.squeeze(axis=(0, 1))
# Flatten all output (0, 1) and input dimensions (2, 3)
final_shape_WR = [WR.shape[0] * WR.shape[1], -1]
assert final_shape_WR[0] == lstm.hidden_size * lstm.multiplier
WR = WR.reshape(final_shape_WR)
final_shape_B = final_shape_WR
if lstm.op == 'GRU' and lstm.linear_before_reset:
final_shape_B[0] = lstm.hidden_size * 4
B = B.reshape(final_shape_B)
# Squeeze fake dimension in B
B = B.squeeze(axis=-1)
for blob, port, name in [(WR, 1, 'weights'), (B, 2, 'biases')]:
Op.create_and_connect_input_data_node(
graph,
lstm,
{'value': blob, 'shape': np.array(blob.shape, dtype=np.int64)},
{'in': port, 'bin': name, 'permutation': None}
)
def repack_weights(graph: Graph, match: dict):
"""
Repack weights into general format (described above) and reorder gates.
"""
rnn_layer = match['rnn_layer']
W = match['W'].value.copy()
R = match['R'].value.copy()
num_directions = 2 if rnn_layer.direction == 'bidirectional' else 1
graph.remove_edge(match['W'].id, rnn_layer.id)
graph.remove_edge(match['R'].id, rnn_layer.id)
# find optional 'B' biases blob
if 3 in rnn_layer.in_nodes():
# TODO: check if 'bin': 'B' attribute is assigned to this edge
B = rnn_layer.in_node(3).value.copy()
graph.remove_edge(rnn_layer.in_node(3).id, rnn_layer.id)
else:
B_shape = [num_directions, 2 * rnn_layer.multiplier * rnn_layer.hidden_size] # from ONNX spec
B = np.full(B_shape, 0, dtype=np.float32)
# Add extra dimensions for W, R and B for easier repacking and reordering
B = B.reshape([
num_directions, # 0: num of directions
rnn_layer.num_layers, # 1: num_layers
2, # 2: two input parts of the matrix: W, R
rnn_layer.multiplier, # 3: four output parts of the matrix for all gates in order: i, o, f, c
rnn_layer.hidden_size, # 4: output size per direction and gate
])
W, R = [x.reshape([
num_directions, # 0: num of directions
rnn_layer.num_layers, # 1: num_layers
rnn_layer.multiplier, # 2: four output parts of the matrix for all gates in order: i, o, f, c
rnn_layer.hidden_size, # 3: output size per direction and gate
-1]) # 4: input size/hidden size in W/R correspondingly
for x in (W, R)]
input_size = match['input'].shape[2]
assert compatible_dims(input_size, W.shape[-1])
# Reorder gates: iofc --> fico
gate_reorder = rnn_layer.gate_order
W, R = (np.take(x, gate_reorder, axis=2) for x in (W, R))
B = np.take(B, gate_reorder, axis=3)
for blob, port in [(W, 1), (R, 2), (B, 3)]:
Op.create_and_connect_input_data_node(
graph,
rnn_layer,
{'value': blob, 'shape': int64_array(blob.shape)},
{'in': port, 'permutation': None}
)
def copy_input_blobs(op: Node, copy_op: Node):
"""
Function copy input blob data nodes from restored graph to copied one
:param op: Node from restored graph
:param copy_op: Node from copied graph
:return:
"""
for u, d in op.get_sorted_inputs():
if 'bin' in d:
Op.create_and_connect_input_data_node(
copy_op.graph, copy_op, {
'value': op.in_node(d['in']).value,
'shape': op.in_node(d['in']).shape
}, d)
def check_init_states(graph: Graph, match: dict):
"""
Check if cell have initial states and create zeros states if not.
And renumber ports for this states.
"""
rnn_cell = match['rnn_layer']
num_directions = 2 if rnn_cell.direction == 'bidirectional' else 1
batch_size = rnn_cell.in_node(0).shape[rnn_cell.batch_dim]
h_init_port = 5
c_init_port = 6
if 2 not in rnn_cell.in_nodes():
h_shape = [num_directions, batch_size,
rnn_cell.hidden_size] # from ONNX spec
h_init = np.full(h_shape, 0, dtype=np.float32)
Op.create_and_connect_input_data_node(
graph, rnn_cell, {
'value': h_init,
'shape': int64_array(h_init.shape)
}, {
'in': h_init_port,
'permutation': None
})
else:
hidden_state_edge = graph.get_edge_data(
rnn_cell.in_node(2).id, rnn_cell.id)
hidden_state_edge[0]['in'] = h_init_port
if rnn_cell.op == 'LSTM':
if 3 not in rnn_cell.in_nodes():
c_shape = [num_directions, batch_size,
rnn_cell.hidden_size] # from ONNX spec
c_init = np.full(c_shape, 0, dtype=np.float32)
Op.create_and_connect_input_data_node(
graph, rnn_cell, {
'value': c_init,
'shape': int64_array(c_init.shape)
}, {
'in': c_init_port,
'permutation': None
})
else:
cell_state_edge = graph.get_edge_data(
rnn_cell.in_node(3).id, rnn_cell.id)
cell_state_edge[0]['in'] = c_init_port
def repack_weights(graph: Graph, match: dict):
input = match['input']
rnn_layer = match['rnn_layer']
params = match['params'].value.copy()
graph.remove_edge(match['params'].id, rnn_layer.id)
input_size = input.shape[2]
direction = 2 if rnn_layer.has_num_directions else 1
bsize = (2 * rnn_layer.hidden_size * direction *
1) * rnn_layer.multiplier
W = mo_array(params[0:len(params) - bsize])
B = mo_array(params[len(params) - bsize:])
W = W.reshape((direction, -1))
B = B.reshape((direction, -1))
W, R = mo_array(W[:, 0:rnn_layer.hidden_size * rnn_layer.multiplier *
input_size]), mo_array(
W[:, rnn_layer.hidden_size *
rnn_layer.multiplier * input_size:])
W, R = [
x.reshape([
direction, # 0: num of directions
1, # 1: num_cells
rnn_layer.
multiplier, # 2: four output parts of the matrix for all gates
rnn_layer.hidden_size, # 3: output size per direction and gate
-1
]) # 4: input size/hidden size in W/R correspondingly
for x in (W, R)
]
assert W.shape[-1] == input_size
assert R.shape[-1] == rnn_layer.hidden_size
B = B.reshape([
direction, # 0: num of directions, limitation: should be 1
1,
2, # 3: num of component B
rnn_layer.
multiplier, # 1: four output parts of the matrix for all gates in order: i, f, c, o
rnn_layer.hidden_size, # 2: output size per direction and gate
])
# Reorder gates: ifco --> fico
gate_reorder = rnn_layer.gate_order
W = np.take(W, gate_reorder, axis=2)
R = np.take(R, gate_reorder, axis=2)
B = np.take(B, gate_reorder, axis=3)
# Add ports to rnn_layer
rnn_layer.add_sequence_of_ports(type='in', rng=range(7))
for blob, port in [(W, 1), (R, 2), (B, 3)]:
Op.create_and_connect_input_data_node(
graph, rnn_layer, {
'value': blob,
'shape': int64_array(blob.shape)
}, {
'in': port,
'permutation': None
})