Roman Kazantsev
b409ea1930
* [TF FE] Support TF1 While Control flow Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com> * Apply code-style fix * Update API for OpPlace to store back edge * Fix build: no rvalue by reference passing * Fix build issue: correct type * Fix TF FE unit-tests * Apply code-review feedback: remove unused vars * Fix fusing complicated case of TF1 While * Remove unused variable * Update MO unit test * Fix layer tests for While * Handle Switch and NextIteration nodes connected directly Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com> --------- Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>
121 lines
5.2 KiB
Python
121 lines
5.2 KiB
Python
# Copyright (C) 2018-2023 Intel Corporation
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import tensorflow as tf
|
|
from common.tf_layer_test_class import CommonTFLayerTest
|
|
|
|
|
|
class TestWhile(CommonTFLayerTest):
|
|
def _prepare_input(self, inputs_info):
|
|
assert 'x' in inputs_info, "Test error: inputs_info must contain `x`"
|
|
assert 'y' in inputs_info, "Test error: inputs_info must contain `y`"
|
|
x_shape = inputs_info['x']
|
|
y_shape = inputs_info['y']
|
|
inputs_data = {}
|
|
inputs_data['x'] = np.random.randint(1, 10, x_shape).astype(np.int32)
|
|
inputs_data['y'] = np.random.randint(-50, 50, y_shape).astype(np.int32)
|
|
return inputs_data
|
|
|
|
def create_while_net(self, y_shape, data_type, lower_control_flow):
|
|
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
|
|
def while_function(x, y):
|
|
@tf.function
|
|
def cond(x, y):
|
|
return tf.less(x, 10)
|
|
|
|
@tf.function
|
|
def body(x, y):
|
|
y_new = tf.add(y, tf.constant(2, dtype=data_type))
|
|
x_new = tf.add(x, 1)
|
|
return x_new, y_new
|
|
|
|
return tf.while_loop(cond, body, [x, y])
|
|
|
|
tf_while_graph = tf.function(while_function)
|
|
x = np.random.randint(1, 10, []).astype(data_type)
|
|
y = np.random.randint(-50, 50, y_shape).astype(data_type)
|
|
concrete_func = tf_while_graph.get_concrete_function(x, y)
|
|
|
|
# lower_control_flow defines representation of While operation
|
|
# in case of lower_control_flow=True it is decomposed into LoopCond, NextIteration and TensorArray operations
|
|
frozen_func = convert_variables_to_constants_v2(concrete_func,
|
|
lower_control_flow=lower_control_flow)
|
|
|
|
graph_def = frozen_func.graph.as_graph_def(add_shapes=True)
|
|
return graph_def, None
|
|
|
|
test_data_basic = [
|
|
dict(y_shape=[2, 3], data_type=np.int32, lower_control_flow=False),
|
|
dict(y_shape=[2, 1, 4], data_type=np.int32, lower_control_flow=False),
|
|
dict(y_shape=[2, 1, 4], data_type=np.int32, lower_control_flow=True)
|
|
]
|
|
|
|
@pytest.mark.parametrize("params", test_data_basic)
|
|
@pytest.mark.precommit_tf_fe
|
|
@pytest.mark.nightly
|
|
def test_while_basic(self, params, ie_device, precision, ir_version, temp_dir,
|
|
use_new_frontend, use_old_api):
|
|
self._test(*self.create_while_net(**params),
|
|
ie_device, precision, ir_version, temp_dir=temp_dir,
|
|
use_new_frontend=use_new_frontend, use_old_api=use_old_api)
|
|
|
|
|
|
class TestWhileShapeVariant(CommonTFLayerTest):
|
|
def _prepare_input(self, inputs_info):
|
|
assert 'x' in inputs_info, "Test error: inputs_info must contain `x`"
|
|
assert 'y' in inputs_info, "Test error: inputs_info must contain `y`"
|
|
x_shape = inputs_info['x']
|
|
y_shape = inputs_info['y']
|
|
inputs_data = {}
|
|
inputs_data['x'] = np.random.randint(1, 10, x_shape).astype(np.int32)
|
|
inputs_data['y'] = np.random.randint(-50, 50, y_shape).astype(np.float32)
|
|
return inputs_data
|
|
|
|
def create_while_net(self, y_shape, lower_control_flow):
|
|
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
|
|
def while_function(x, y):
|
|
@tf.function
|
|
def cond(x, y):
|
|
return tf.less(x, 10)
|
|
|
|
@tf.function
|
|
def body(x, y):
|
|
add_2 = tf.add(y, tf.constant(2, dtype=tf.float32))
|
|
y_new = tf.concat(values=[y, add_2], axis=0)
|
|
x_new = tf.add(x, tf.constant(1, tf.int32))
|
|
return x_new, y_new
|
|
|
|
return tf.while_loop(cond, body, [x, y],
|
|
shape_invariants=[tf.TensorShape([]),
|
|
tf.TensorShape([None] + y_shape[1:])])
|
|
|
|
tf_while_graph = tf.function(while_function)
|
|
x = np.random.randint(1, 10, []).astype(np.int32)
|
|
y = np.random.randint(-50, 50, y_shape).astype(np.float32)
|
|
concrete_func = tf_while_graph.get_concrete_function(x, y)
|
|
|
|
# lower_control_flow defines representation of While operation
|
|
# in case of lower_control_flow=True it is decomposed into LoopCond, NextIteration and TensorArray operations
|
|
frozen_func = convert_variables_to_constants_v2(concrete_func,
|
|
lower_control_flow=lower_control_flow)
|
|
|
|
graph_def = frozen_func.graph.as_graph_def(add_shapes=True)
|
|
return graph_def, None
|
|
|
|
test_data_basic = [
|
|
dict(y_shape=[2, 3], lower_control_flow=False),
|
|
dict(y_shape=[2, 1, 4], lower_control_flow=False),
|
|
dict(y_shape=[2, 1, 4], lower_control_flow=True)
|
|
]
|
|
|
|
@pytest.mark.parametrize("params", test_data_basic)
|
|
@pytest.mark.precommit_tf_fe
|
|
@pytest.mark.nightly
|
|
def test_while_basic(self, params, ie_device, precision, ir_version, temp_dir,
|
|
use_new_frontend, use_old_api):
|
|
self._test(*self.create_while_net(**params),
|
|
ie_device, precision, ir_version, temp_dir=temp_dir,
|
|
use_new_frontend=use_new_frontend, use_old_api=use_old_api)
|