• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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