1# Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Test configs for unidirectional_sequence_rnn.""" 16from __future__ import absolute_import 17from __future__ import division 18from __future__ import print_function 19 20import tensorflow.compat.v1 as tf 21from tensorflow.lite.testing.zip_test_utils import create_tensor_data 22from tensorflow.lite.testing.zip_test_utils import make_zip_of_tests 23from tensorflow.lite.testing.zip_test_utils import register_make_test_function 24from tensorflow.python.framework import test_util 25 26 27@register_make_test_function("make_unidirectional_sequence_rnn_tests") 28@test_util.enable_control_flow_v2 29def make_unidirectional_sequence_rnn_tests(options): 30 """Make a set of tests to do unidirectional_sequence_rnn.""" 31 32 test_parameters = [{ 33 "batch_size": [2, 4, 6], 34 "seq_length": [1, 3], 35 "units": [4, 5], 36 "is_dynamic_rnn": [False, True] 37 }] 38 39 def build_graph(parameters): 40 """Build the graph for unidirectional_sequence_rnn.""" 41 input_values = [] 42 if parameters["is_dynamic_rnn"]: 43 shape = [ 44 parameters["seq_length"], parameters["batch_size"], 45 parameters["units"] 46 ] 47 input_value = tf.compat.v1.placeholder( 48 dtype=tf.float32, name="input", shape=shape) 49 input_values.append(input_value) 50 rnn_cell = tf.lite.experimental.nn.TfLiteRNNCell(parameters["units"]) 51 outs, _ = tf.lite.experimental.nn.dynamic_rnn( 52 rnn_cell, input_value, dtype=tf.float32, time_major=True) 53 outs = tf.unstack(outs, axis=1) 54 else: 55 shape = [parameters["batch_size"], parameters["units"]] 56 for i in range(parameters["seq_length"]): 57 input_value = tf.compat.v1.placeholder( 58 dtype=tf.float32, name=("input_%d" % i), shape=shape) 59 input_values.append(input_value) 60 rnn_cell = tf.lite.experimental.nn.TfLiteRNNCell(parameters["units"]) 61 outs, _ = tf.nn.static_rnn(rnn_cell, input_values, dtype=tf.float32) 62 63 real_output = tf.zeros([1], dtype=tf.float32) + outs[-1] 64 real_output = tf.identity(real_output) 65 return input_values, [real_output] 66 67 def build_inputs(parameters, sess, inputs, outputs): 68 """Build the inputs for unidirectional_sequence_rnn.""" 69 input_values = [] 70 if parameters["is_dynamic_rnn"]: 71 shape = [ 72 parameters["seq_length"], parameters["batch_size"], 73 parameters["units"] 74 ] 75 input_value = create_tensor_data(tf.float32, shape) 76 input_values.append(input_value) 77 else: 78 shape = [parameters["batch_size"], parameters["units"]] 79 for _ in range(parameters["seq_length"]): 80 input_value = create_tensor_data(tf.float32, shape) 81 input_values.append(input_value) 82 init = tf.compat.v1.global_variables_initializer() 83 sess.run(init) 84 # Tflite fused kernel takes input as [time, batch, input]. 85 # For static unidirectional sequence rnn, the input is an array sized of 86 # time, and pack the array together, however, for time = 1, the input is 87 # not packed. 88 tflite_input_values = input_values 89 if not parameters["is_dynamic_rnn"] and parameters["seq_length"] == 1: 90 tflite_input_values = [ 91 input_values[0].reshape( 92 (1, parameters["batch_size"], parameters["units"])) 93 ] 94 return tflite_input_values, sess.run( 95 outputs, feed_dict=dict(zip(inputs, input_values))) 96 97 make_zip_of_tests( 98 options, 99 test_parameters, 100 build_graph, 101 build_inputs, 102 use_frozen_graph=True) 103