64 lines
2.7 KiB
Python
64 lines
2.7 KiB
Python
"""
|
|
Copyright (C) 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.
|
|
"""
|
|
|
|
import numpy as np
|
|
|
|
from mo.graph.graph import Graph
|
|
from mo.middle.replacement import MiddleReplacementPattern
|
|
from mo.ops.const import Const
|
|
from mo.utils.error import Error
|
|
|
|
|
|
class SequenceLengthToMask(MiddleReplacementPattern):
|
|
"""
|
|
Convert a sequence length to a sequence mask for CTCGreedyDecoder if its value is available.
|
|
"""
|
|
enabled = True
|
|
|
|
def run_before(self):
|
|
from extensions.middle.pass_separator import MiddleFinish
|
|
return [MiddleFinish]
|
|
|
|
def find_and_replace_pattern(self, graph: Graph):
|
|
for ctc_greedy_decoder in graph.get_op_nodes(op='CTCGreedyDecoder', use_mask_format=True):
|
|
ctc_greedy_decoder_name = ctc_greedy_decoder.soft_get('name', ctc_greedy_decoder.id)
|
|
|
|
sequence_length_value = ctc_greedy_decoder.in_port(1).data.get_value()
|
|
if sequence_length_value is None:
|
|
raise Error('The second input to the CTCGreedyDecoder node "{}" is not constant. This case is not '
|
|
'supported with the Inference Engine.'.format(ctc_greedy_decoder_name))
|
|
|
|
# transform a sequence length to a sequence mask
|
|
logits_shape = ctc_greedy_decoder.in_port(0).data.get_shape()
|
|
assert logits_shape is not None and len(logits_shape) == 3, \
|
|
"Incorrect shape for logits input of {} node".format(ctc_greedy_decoder_name)
|
|
batch_size = logits_shape[1]
|
|
time_size = logits_shape[0]
|
|
mask_value = np.zeros([batch_size, time_size], dtype=np.float)
|
|
for sample_ind, sample_seq_length in enumerate(sequence_length_value):
|
|
mask_value[sample_ind, 0:sample_seq_length] = 1
|
|
mask_value = np.transpose(mask_value)
|
|
|
|
# create Const node with computed mask value
|
|
mask_node = Const(graph, {'name': ctc_greedy_decoder_name + '/Mask',
|
|
'value': mask_value}).create_node()
|
|
|
|
# connect computed mask to CTCGreedyDecoder node
|
|
ctc_greedy_decoder.in_port(1).get_connection().set_source(mask_node.out_port(0))
|
|
|
|
# remove attribute-marker
|
|
del ctc_greedy_decoder['use_mask_format']
|