openvino/model-optimizer/mo/front/kaldi/extractors/timeheightconvolution_ext.py

63 lines
3.0 KiB
Python

# Copyright (C) 2021 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import numpy as np
from mo.front.caffe.extractors.utils import embed_input
from mo.front.extractor import FrontExtractorOp
from mo.front.kaldi.loader.utils import collect_until_token, read_token_value
from mo.front.kaldi.utils import read_binary_matrix, read_binary_vector, read_binary_vector_of_pairs
from mo.ops.timeheightconvolution import TimeHeightConvolutionComponent
class TimeHeightConvolutionFrontExtractor(FrontExtractorOp):
op = 'timeheightconvolutioncomponent'
enabled = True
@classmethod
def extract(cls, node):
pb = node.parameters
collect_until_token(pb, b'<ConvolutionModel>')
in_shape = read_token_value(pb, b'<NumFiltersIn>')
out_shape = read_token_value(pb, b'<NumFiltersOut>')
height_in = read_token_value(pb, b'<HeightIn>')
height_out = read_token_value(pb, b'<HeightOut>')
height_subsample = read_token_value(pb, b'<HeightSubsampleOut>')
collect_until_token(pb, b'<Offsets>')
offsets = read_binary_vector_of_pairs(pb, read_token=False, dtype=np.int32)
collect_until_token(pb, b'<RequiredTimeOffsets>')
time_offsets = read_binary_vector(pb, read_token=False, dtype=np.int32)
collect_until_token(pb, b'<LinearParams>')
weights, _ = read_binary_matrix(pb)
collect_until_token(pb, b'<BiasParams>')
biases = read_binary_vector(pb)
offsets = offsets.reshape([len(offsets)//2, 2])
mapping_rule = { # stride for h axis
'height_subsample': height_subsample,
# input dimension for h axis
'height_in': height_in,
# output dimension for h axis
'height_out': height_out,
# input dimension for channel axis
'in_channels': in_shape,
# output dimension for channel axis
'out_channels': out_shape,
# array with pairs like the following
# [ (-1, -1) (-1, 0) (-1, 1)
# (0, -1) (0, 0) (0, 1)
# (1, -1) (1, 0) (1, 1)]
# it means that kernel 3x3 will be applied to calculate current value of output
'offsets': offsets,
# required time offsets to calculate current convolution
# time_offsets = [-1, 0, 1] for previous example means no padding for time axis and
# 3 values should be prepared
# time_offsets = [0] means zero padding [1, 1] for time axis
'time_offsets': time_offsets,
'out-size': out_shape * height_out}
embed_input(mapping_rule, 1, 'weights', weights)
embed_input(mapping_rule, 2, 'biases', biases)
TimeHeightConvolutionComponent.update_node_stat(node, mapping_rule)
return cls.enabled