59 lines
2.4 KiB
Python
59 lines
2.4 KiB
Python
# Copyright (C) 2018-2021 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
|
|
from extensions.front.mxnet.gather import GatherFrontReplacer
|
|
from mo.utils.ir_engine.compare_graphs import compare_graphs
|
|
from unit_tests.utils.graph import build_graph
|
|
|
|
|
|
class GatherTest(unittest.TestCase):
|
|
def test_embedding_replace1(self):
|
|
graph = build_graph(
|
|
{'placeholder_1': {'shape': None, 'type': 'Parameter', 'kind': 'op', 'op': 'Parameter'},
|
|
'embedding_const': {'value': None, 'shape': None, 'kind': 'op', 'op': 'Const'},
|
|
'embedding': {'type': None, 'kind': 'op', 'op': 'Embedding'},
|
|
'last': {'type': None, 'kind': 'op', 'op': None},
|
|
},
|
|
[
|
|
('placeholder_1', 'embedding', {'out': 0, 'in': 0}),
|
|
('embedding_const', 'embedding', {'out': 0, 'in': 1}),
|
|
('embedding', 'last')
|
|
],
|
|
{
|
|
'placeholder_1': {'shape': np.array([32, 35])},
|
|
'embedding_const': {'shape': np.array([2000, 650]),
|
|
'bias': np.array(np.random.randint(0, 225, (2000, 650)))},
|
|
}, nodes_with_edges_only=True)
|
|
|
|
graph_ref = build_graph(
|
|
{'placeholder_1': {'shape': None, 'type': 'Parameter', 'kind': 'op', 'op': 'Parameter'},
|
|
'embedding_const': {'value': None, 'kind': 'op', 'op': 'Const'},
|
|
'axis_const': {'value': 0, 'kind': 'op', 'data_type': None,
|
|
'type': 'Const', 'op': 'Const'},
|
|
'embedding': {'kind': 'op', 'op': 'Gather'},
|
|
'last': {'type': None, 'kind': 'op', 'op': None},
|
|
},
|
|
[
|
|
('embedding_const', 'embedding', {'in': 1}),
|
|
('axis_const', 'embedding', {'in': 2}),
|
|
('placeholder_1', 'embedding', {'in': 0}),
|
|
('embedding', 'last')
|
|
],
|
|
{'placeholder_1': {'shape': np.array([32, 35])},
|
|
'embedding_const': {'shape': np.array([2000, 650]),
|
|
'bias': np.array(np.random.randint(0, 225, (2000, 650)))},
|
|
}, nodes_with_edges_only=True)
|
|
|
|
graph.graph['layout'] = 'NCHW'
|
|
graph.stage = 'front'
|
|
|
|
replacer = GatherFrontReplacer()
|
|
replacer.find_and_replace_pattern(graph)
|
|
|
|
(flag, resp) = compare_graphs(graph, graph_ref, 'last')
|
|
self.assertTrue(flag, resp)
|