1# Copyright 2018 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"""Tests for cudnn recurrent layers."""
16
17from __future__ import absolute_import
18from __future__ import division
19from __future__ import print_function
20
21import os
22import tempfile
23from absl.testing import parameterized
24import numpy as np
25
26from tensorflow.python import keras
27from tensorflow.python.framework import test_util
28from tensorflow.python.keras import keras_parameterized
29from tensorflow.python.keras import testing_utils
30from tensorflow.python.keras.optimizer_v2.rmsprop import RMSprop
31from tensorflow.python.ops import array_ops
32from tensorflow.python.platform import test
33from tensorflow.python.training import gradient_descent
34
35
36@keras_parameterized.run_all_keras_modes
37class CuDNNTest(keras_parameterized.TestCase):
38
39  @parameterized.named_parameters(
40      *test_util.generate_combinations_with_testcase_name(
41          layer_class=[keras.layers.CuDNNGRU, keras.layers.CuDNNLSTM],
42          return_sequences=[True, False]))
43  @test_util.run_gpu_only
44  def test_cudnn_rnn_return_sequence(self, layer_class, return_sequences):
45    input_size = 10
46    timesteps = 6
47    units = 2
48    num_samples = 32
49    testing_utils.layer_test(
50        layer_class,
51        kwargs={'units': units,
52                'return_sequences': return_sequences},
53        input_shape=(num_samples, timesteps, input_size))
54
55  @parameterized.named_parameters(
56      *test_util.generate_combinations_with_testcase_name(
57          layer_class=[keras.layers.CuDNNGRU, keras.layers.CuDNNLSTM],
58          go_backwards=[True, False]))
59  @test_util.run_gpu_only
60  def test_cudnn_rnn_go_backward(self, layer_class, go_backwards):
61    input_size = 10
62    timesteps = 6
63    units = 2
64    num_samples = 32
65    testing_utils.layer_test(
66        layer_class,
67        kwargs={'units': units,
68                'go_backwards': go_backwards},
69        input_shape=(num_samples, timesteps, input_size))
70
71  @parameterized.named_parameters(
72      ('cudnngru', keras.layers.CuDNNGRU),
73      ('cudnnlstm', keras.layers.CuDNNLSTM),
74  )
75  @test_util.run_gpu_only
76  def test_return_state(self, layer_class):
77    input_size = 10
78    timesteps = 6
79    units = 2
80    num_samples = 32
81    num_states = 2 if layer_class is keras.layers.CuDNNLSTM else 1
82
83    inputs = keras.Input(batch_shape=(num_samples, timesteps, input_size))
84    layer = layer_class(units, return_state=True, stateful=True)
85    outputs = layer(inputs)
86    _, state = outputs[0], outputs[1:]
87    self.assertEqual(len(state), num_states)
88    model = keras.models.Model(inputs, state[0])
89    model.run_eagerly = testing_utils.should_run_eagerly()
90
91    inputs = np.random.random((num_samples, timesteps, input_size))
92    state = model.predict(inputs)
93    np.testing.assert_allclose(
94        keras.backend.eval(layer.states[0]), state, atol=1e-4)
95
96  @parameterized.named_parameters(
97      ('cudnngru', keras.layers.CuDNNGRU),
98      ('cudnnlstm', keras.layers.CuDNNLSTM),
99  )
100  @test_util.run_gpu_only
101  def test_time_major_input(self, layer_class):
102    input_size = 10
103    timesteps = 6
104    units = 2
105    num_samples = 32
106
107    model = keras.models.Sequential()
108    model.add(
109        keras.layers.Lambda(lambda t: array_ops.transpose(t, [1, 0, 2])))
110    layer = layer_class(units, time_major=True, return_sequences=True)
111    model.add(layer)
112    model.add(
113        keras.layers.Lambda(lambda t: array_ops.transpose(t, [1, 0, 2])))
114    model.compile(loss='categorical_crossentropy',
115                  optimizer=RMSprop(learning_rate=0.001))
116    model.fit(
117        np.ones((num_samples, timesteps, input_size)),
118        np.ones((num_samples, timesteps, units)))
119    out = model.predict(np.ones((num_samples, timesteps, input_size)))
120    self.assertEqual(out.shape, (num_samples, timesteps, units))
121
122  @parameterized.named_parameters(
123      ('cudnngru', keras.layers.CuDNNGRU),
124      ('cudnnlstm', keras.layers.CuDNNLSTM),
125  )
126  @test_util.run_gpu_only
127  def test_specify_initial_state_keras_tensor(self, layer_class):
128    input_size = 10
129    timesteps = 6
130    units = 2
131    num_samples = 32
132    num_states = 2 if layer_class is keras.layers.CuDNNLSTM else 1
133
134    inputs = keras.Input((timesteps, input_size))
135    initial_state = [keras.Input((units,)) for _ in range(num_states)]
136    layer = layer_class(units)
137    if len(initial_state) == 1:
138      output = layer(inputs, initial_state=initial_state[0])
139    else:
140      output = layer(inputs, initial_state=initial_state)
141    self.assertIn(initial_state[0], layer._inbound_nodes[0].input_tensors)
142
143    model = keras.models.Model([inputs] + initial_state, output)
144    model.compile(loss='categorical_crossentropy',
145                  optimizer=RMSprop(learning_rate=0.001),
146                  run_eagerly=testing_utils.should_run_eagerly())
147
148    inputs = np.random.random((num_samples, timesteps, input_size))
149    initial_state = [
150        np.random.random((num_samples, units)) for _ in range(num_states)
151    ]
152    targets = np.random.random((num_samples, units))
153    model.fit([inputs] + initial_state, targets)
154
155
156class CuDNNGraphOnlyTest(keras_parameterized.TestCase):
157
158  @parameterized.named_parameters(
159      ('cudnngru', keras.layers.CuDNNGRU),
160      ('cudnnlstm', keras.layers.CuDNNLSTM),
161  )
162  @test_util.run_deprecated_v1
163  @test_util.run_gpu_only
164  def test_regularizer(self, layer_class):
165    input_size = 10
166    timesteps = 6
167    units = 2
168    num_samples = 32
169    layer = layer_class(
170        units,
171        return_sequences=False,
172        input_shape=(timesteps, input_size),
173        kernel_regularizer=keras.regularizers.l1(0.01),
174        recurrent_regularizer=keras.regularizers.l1(0.01),
175        bias_regularizer='l2')
176    layer.build((None, None, input_size))
177    self.assertEqual(len(layer.losses), 3)
178
179    layer = layer_class(
180        units,
181        return_sequences=False,
182        input_shape=(timesteps, input_size),
183        activity_regularizer='l2')
184    self.assertTrue(layer.activity_regularizer)
185    x = keras.backend.variable(
186        np.ones((num_samples, timesteps, input_size)))
187    layer(x)
188    self.assertEqual(len(layer.get_losses_for(x)), 1)
189
190  @parameterized.named_parameters(
191      ('cudnngru', keras.layers.CuDNNGRU),
192      ('cudnnlstm', keras.layers.CuDNNLSTM),
193  )
194  @test_util.run_gpu_only
195  @test_util.run_v1_only('b/120941292')
196  def test_statefulness(self, layer_class):
197    input_size = 10
198    timesteps = 6
199    units = 2
200    num_samples = 32
201
202    with self.cached_session(use_gpu=True):
203      model = keras.models.Sequential()
204      model.add(
205          keras.layers.Embedding(
206              10,
207              input_size,
208              input_length=timesteps,
209              batch_input_shape=(num_samples, timesteps)))
210      layer = layer_class(
211          units, return_sequences=False, stateful=True, weights=None)
212      model.add(layer)
213      model.compile(optimizer=gradient_descent.GradientDescentOptimizer(0.01),
214                    loss='mse')
215      out1 = model.predict(np.ones((num_samples, timesteps)))
216      self.assertEqual(out1.shape, (num_samples, units))
217
218      # train once so that the states change
219      model.train_on_batch(
220          np.ones((num_samples, timesteps)), np.ones((num_samples, units)))
221      out2 = model.predict(np.ones((num_samples, timesteps)))
222
223      # if the state is not reset, output should be different
224      self.assertNotEqual(out1.max(), out2.max())
225
226      # check that output changes after states are reset
227      # (even though the model itself didn't change)
228      layer.reset_states()
229      out3 = model.predict(np.ones((num_samples, timesteps)))
230      self.assertNotEqual(out2.max(), out3.max())
231
232      # check that container-level reset_states() works
233      model.reset_states()
234      out4 = model.predict(np.ones((num_samples, timesteps)))
235      self.assertAllClose(out3, out4, atol=1e-5)
236
237      # check that the call to `predict` updated the states
238      out5 = model.predict(np.ones((num_samples, timesteps)))
239      self.assertNotEqual(out4.max(), out5.max())
240
241
242@test_util.run_all_in_graph_and_eager_modes
243class CuDNNV1OnlyTest(keras_parameterized.TestCase):
244
245  @test_util.run_gpu_only
246  def test_trainability(self):
247    input_size = 10
248    units = 2
249    for layer_class in [keras.layers.CuDNNGRU, keras.layers.CuDNNLSTM]:
250      layer = layer_class(units)
251      layer.build((None, None, input_size))
252      self.assertEqual(len(layer.weights), 3)
253      self.assertEqual(len(layer.trainable_weights), 3)
254      self.assertEqual(len(layer.non_trainable_weights), 0)
255      layer.trainable = False
256      self.assertEqual(len(layer.weights), 3)
257      self.assertEqual(len(layer.non_trainable_weights), 3)
258      self.assertEqual(len(layer.trainable_weights), 0)
259      layer.trainable = True
260      self.assertEqual(len(layer.weights), 3)
261      self.assertEqual(len(layer.trainable_weights), 3)
262      self.assertEqual(len(layer.non_trainable_weights), 0)
263
264  @parameterized.named_parameters(
265      *test_util.generate_combinations_with_testcase_name(
266          rnn_type=['LSTM', 'GRU'], to_cudnn=[True, False],
267          bidirectional=[True, False], implementation=[1, 2],
268          model_nest_level=[1, 2], model_type=['seq', 'func']))
269  @test_util.run_v1_only('b/120911602, b/112083752')
270  @test_util.run_gpu_only
271  def test_load_weights_between_noncudnn_rnn(self, rnn_type, to_cudnn,
272                                             bidirectional, implementation,
273                                             model_nest_level, model_type):
274    input_size = 10
275    timesteps = 6
276    input_shape = (timesteps, input_size)
277    units = 2
278    num_samples = 32
279    inputs = np.random.random((num_samples, timesteps, input_size))
280
281    rnn_layer_kwargs = {
282        'recurrent_activation': 'sigmoid',
283        # ensure biases are non-zero and properly converted
284        'bias_initializer': 'random_uniform',
285        'implementation': implementation
286    }
287    if rnn_type == 'LSTM':
288      rnn_layer_class = keras.layers.LSTM
289      cudnn_rnn_layer_class = keras.layers.CuDNNLSTM
290    else:
291      rnn_layer_class = keras.layers.GRU
292      cudnn_rnn_layer_class = keras.layers.CuDNNGRU
293      rnn_layer_kwargs['reset_after'] = True
294
295    layer = rnn_layer_class(units, **rnn_layer_kwargs)
296    if bidirectional:
297      layer = keras.layers.Bidirectional(layer)
298
299    cudnn_layer = cudnn_rnn_layer_class(units)
300    if bidirectional:
301      cudnn_layer = keras.layers.Bidirectional(cudnn_layer)
302
303    model = self._make_nested_model(input_shape, layer, model_nest_level,
304                                    model_type)
305    cudnn_model = self._make_nested_model(input_shape, cudnn_layer,
306                                          model_nest_level, model_type)
307
308    if to_cudnn:
309      self._convert_model_weights(model, cudnn_model)
310    else:
311      self._convert_model_weights(cudnn_model, model)
312
313    self.assertAllClose(model.predict(inputs), cudnn_model.predict(inputs),
314                        atol=1e-4)
315
316  def _make_nested_model(self, input_shape, layer, level=1, model_type='func'):
317    # example: make_nested_seq_model((1,), Dense(10), level=2).summary()
318    def make_nested_seq_model(input_shape, layer, level=1):
319      model = layer
320      for i in range(1, level + 1):
321        layers = [keras.layers.InputLayer(input_shape),
322                  model] if (i == 1) else [model]
323        model = keras.models.Sequential(layers)
324      return model
325
326    # example: make_nested_func_model((1,), Dense(10), level=2).summary()
327    def make_nested_func_model(input_shape, layer, level=1):
328      model_input = keras.layers.Input(input_shape)
329      model = layer
330      for _ in range(level):
331        model = keras.models.Model(model_input, model(model_input))
332      return model
333
334    if model_type == 'func':
335      return make_nested_func_model(input_shape, layer, level)
336    elif model_type == 'seq':
337      return make_nested_seq_model(input_shape, layer, level)
338
339  def _convert_model_weights(self, source_model, target_model):
340    _, fname = tempfile.mkstemp('.h5')
341    source_model.save_weights(fname)
342    target_model.load_weights(fname)
343    os.remove(fname)
344
345  @parameterized.named_parameters(
346      *test_util.generate_combinations_with_testcase_name(
347          rnn_type=['LSTM', 'GRU'], to_cudnn=[True, False]))
348  @test_util.run_v1_only('b/120911602')
349  @test_util.run_gpu_only
350  def test_load_weights_between_noncudnn_rnn_time_distributed(self, rnn_type,
351                                                              to_cudnn):
352    # Similar test as test_load_weights_between_noncudnn_rnn() but has different
353    # rank of input due to usage of TimeDistributed. Issue: #10356.
354    input_size = 10
355    steps = 6
356    timesteps = 6
357    input_shape = (timesteps, steps, input_size)
358    units = 2
359    num_samples = 32
360    inputs = np.random.random((num_samples, timesteps, steps, input_size))
361
362    rnn_layer_kwargs = {
363        'recurrent_activation': 'sigmoid',
364        # ensure biases are non-zero and properly converted
365        'bias_initializer': 'random_uniform',
366    }
367    if rnn_type == 'LSTM':
368      rnn_layer_class = keras.layers.LSTM
369      cudnn_rnn_layer_class = keras.layers.CuDNNLSTM
370    else:
371      rnn_layer_class = keras.layers.GRU
372      cudnn_rnn_layer_class = keras.layers.CuDNNGRU
373      rnn_layer_kwargs['reset_after'] = True
374
375    layer = rnn_layer_class(units, **rnn_layer_kwargs)
376    layer = keras.layers.TimeDistributed(layer)
377
378    cudnn_layer = cudnn_rnn_layer_class(units)
379    cudnn_layer = keras.layers.TimeDistributed(cudnn_layer)
380
381    model = self._make_nested_model(input_shape, layer)
382    cudnn_model = self._make_nested_model(input_shape, cudnn_layer)
383
384    if to_cudnn:
385      self._convert_model_weights(model, cudnn_model)
386    else:
387      self._convert_model_weights(cudnn_model, model)
388
389    self.assertAllClose(model.predict(inputs), cudnn_model.predict(inputs),
390                        atol=1e-4)
391
392  @test_util.run_gpu_only
393  def test_cudnnrnn_bidirectional(self):
394    rnn = keras.layers.CuDNNGRU
395    samples = 2
396    dim = 2
397    timesteps = 2
398    output_dim = 2
399    mode = 'concat'
400
401    x = np.random.random((samples, timesteps, dim))
402    target_dim = 2 * output_dim if mode == 'concat' else output_dim
403    y = np.random.random((samples, target_dim))
404
405    # test with Sequential model
406    model = keras.Sequential()
407    model.add(
408        keras.layers.Bidirectional(
409            rnn(output_dim), merge_mode=mode, input_shape=(None, dim)))
410    model.compile(loss='mse', optimizer='rmsprop')
411    model.fit(x, y, epochs=1, batch_size=1)
412
413    # test config
414    model.get_config()
415    model = keras.models.model_from_json(model.to_json())
416    model.summary()
417
418    # test stacked bidirectional layers
419    model = keras.Sequential()
420    model.add(
421        keras.layers.Bidirectional(
422            rnn(output_dim, return_sequences=True),
423            merge_mode=mode,
424            input_shape=(None, dim)))
425    model.add(keras.layers.Bidirectional(rnn(output_dim), merge_mode=mode))
426    model.compile(loss='mse', optimizer=R'rmsprop')
427    model.fit(x, y, epochs=1, batch_size=1)
428
429    # test with functional API
430    inputs = keras.Input((timesteps, dim))
431    outputs = keras.layers.Bidirectional(
432        rnn(output_dim), merge_mode=mode)(
433            inputs)
434    model = keras.Model(inputs, outputs)
435    model.compile(loss='mse', optimizer=R'rmsprop')
436    model.fit(x, y, epochs=1, batch_size=1)
437
438    # Bidirectional and stateful
439    inputs = keras.Input(batch_shape=(1, timesteps, dim))
440    outputs = keras.layers.Bidirectional(
441        rnn(output_dim, stateful=True), merge_mode=mode)(
442            inputs)
443    model = keras.Model(inputs, outputs)
444    model.compile(loss='mse', optimizer='rmsprop')
445    model.fit(x, y, epochs=1, batch_size=1)
446
447  @test_util.run_gpu_only
448  def test_preprocess_weights_for_loading_gru_incompatible(self):
449    """Test loading weights between incompatible layers.
450
451    Should fail fast with an exception.
452    """
453    input_shape = (3, 5)
454
455    def gru(cudnn=False, **kwargs):
456      layer_class = keras.layers.CuDNNGRU if cudnn else keras.layers.GRU
457      return layer_class(2, input_shape=input_shape, **kwargs)
458
459    def get_layer_weights(layer):
460      layer.build(input_shape=input_shape)
461      return layer.get_weights()
462
463    def assert_not_compatible(src, dest, message):
464      with self.assertRaises(ValueError) as ex:
465        keras.saving.preprocess_weights_for_loading(
466            dest,
467            get_layer_weights(src))
468      self.assertIn(message, str(ex.exception))
469
470    assert_not_compatible(
471        gru(),
472        gru(cudnn=True),
473        'GRU(reset_after=False) is not compatible with CuDNNGRU')
474    assert_not_compatible(
475        gru(cudnn=True),
476        gru(),
477        'CuDNNGRU is not compatible with GRU(reset_after=False)')
478    assert_not_compatible(
479        gru(),
480        gru(reset_after=True),
481        'GRU(reset_after=False) is not compatible with '
482        'GRU(reset_after=True)')
483    assert_not_compatible(
484        gru(reset_after=True),
485        gru(),
486        'GRU(reset_after=True) is not compatible with '
487        'GRU(reset_after=False)')
488
489
490if __name__ == '__main__':
491  test.main()
492