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 UnifiedLSTM layer.""" 16 17 from __future__ import absolute_import 18 from __future__ import division 19 from __future__ import print_function 20 21 import os 22 import shutil 23 import time 24 25 from absl.testing import parameterized 26 import numpy as np 27 28 from tensorflow.core.protobuf import config_pb2 29 from tensorflow.core.protobuf import rewriter_config_pb2 30 from tensorflow.python import keras 31 from tensorflow.python.client import session as session_lib 32 from tensorflow.python.eager import context 33 from tensorflow.python.framework import constant_op 34 from tensorflow.python.framework import dtypes 35 from tensorflow.python.framework import test_util 36 from tensorflow.python.keras import keras_parameterized 37 from tensorflow.python.keras import testing_utils 38 from tensorflow.python.keras.layers import recurrent as rnn_v1 39 from tensorflow.python.keras.layers import recurrent_v2 as rnn 40 from tensorflow.python.ops import array_ops 41 from tensorflow.python.ops import control_flow_ops 42 from tensorflow.python.ops import gen_math_ops 43 from tensorflow.python.ops import variables 44 from tensorflow.python.ops.losses import losses 45 from tensorflow.python.platform import test 46 from tensorflow.python.platform import tf_logging as logging 47 from tensorflow.python.training import gradient_descent 48 from tensorflow.python.util import nest 49 50 51 # Global config for grappler setting that is used for graph mode test. 52 _rewrites = rewriter_config_pb2.RewriterConfig() 53 _rewrites.implementation_selector = rewriter_config_pb2.RewriterConfig.ON 54 _rewrites.min_graph_nodes = -1 55 _graph_options = config_pb2.GraphOptions(rewrite_options=_rewrites) 56 _config = config_pb2.ConfigProto(graph_options=_graph_options) 57 58 59 @keras_parameterized.run_all_keras_modes(config=_config) 60 class LSTMV2Test(keras_parameterized.TestCase): 61 62 @parameterized.named_parameters( 63 ('non_tan_activation', 'relu', 'sigmoid', 0, False, True), 64 ('non_sigmoid_recur_activation', 'tanh', 'relu', 0, False, True), 65 ('use_recurrent_dropout', 'tanh', 'sigmoid', 0.1, False, True), 66 ('unroll', 'tanh', 'sigmoid', 0, True, True), 67 ('not_use_bias', 'tanh', 'sigmoid', 0, False, False), 68 ) 69 def test_could_use_defun_backend(self, activation, recurrent_activation, 70 recurrent_dropout, unroll, use_bias): 71 layer = rnn.LSTM( 72 1, 73 activation=activation, 74 recurrent_activation=recurrent_activation, 75 recurrent_dropout=recurrent_dropout, 76 unroll=unroll, 77 use_bias=use_bias) 78 self.assertFalse(layer.could_use_cudnn) 79 80 def test_static_shape_inference_LSTM(self): 81 # Github issue: 15165 82 timesteps = 3 83 embedding_dim = 4 84 units = 2 85 86 model = keras.models.Sequential() 87 inputs = keras.layers.Dense( 88 embedding_dim, input_shape=(timesteps, embedding_dim)) 89 model.add(inputs) 90 layer = rnn.LSTM(units, return_sequences=True) 91 model.add(layer) 92 outputs = model.layers[-1].output 93 self.assertEqual(outputs.get_shape().as_list(), [None, timesteps, units]) 94 95 def test_dynamic_behavior_LSTM(self): 96 num_samples = 2 97 timesteps = 3 98 embedding_dim = 4 99 units = 2 100 layer = rnn.LSTM(units, input_shape=(None, embedding_dim)) 101 model = keras.models.Sequential() 102 model.add(layer) 103 model.compile(gradient_descent.GradientDescentOptimizer(0.001), 'mse') 104 x = np.random.random((num_samples, timesteps, embedding_dim)) 105 y = np.random.random((num_samples, units)) 106 model.train_on_batch(x, y) 107 108 def test_stacking_LSTM(self): 109 inputs = np.random.random((2, 3, 4)) 110 targets = np.abs(np.random.random((2, 3, 5))) 111 targets /= targets.sum(axis=-1, keepdims=True) 112 model = keras.models.Sequential() 113 model.add(rnn.LSTM(10, return_sequences=True, unroll=False)) 114 model.add(rnn.LSTM(5, return_sequences=True, unroll=False)) 115 model.compile( 116 loss='categorical_crossentropy', 117 optimizer=gradient_descent.GradientDescentOptimizer(0.01)) 118 model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1) 119 120 def test_from_config_LSTM(self): 121 layer_class = rnn.LSTM 122 for stateful in (False, True): 123 l1 = layer_class(units=1, stateful=stateful) 124 l2 = layer_class.from_config(l1.get_config()) 125 assert l1.get_config() == l2.get_config() 126 127 def test_specify_initial_state_keras_tensor(self): 128 num_states = 2 129 timesteps = 3 130 embedding_dim = 4 131 units = 3 132 num_samples = 2 133 134 # Test with Keras tensor 135 inputs = keras.Input((timesteps, embedding_dim)) 136 initial_state = [keras.Input((units,)) for _ in range(num_states)] 137 layer = rnn.LSTM(units) 138 if len(initial_state) == 1: 139 output = layer(inputs, initial_state=initial_state[0]) 140 else: 141 output = layer(inputs, initial_state=initial_state) 142 assert initial_state[0] in layer._inbound_nodes[0].input_tensors 143 144 model = keras.models.Model([inputs] + initial_state, output) 145 model.compile( 146 loss='categorical_crossentropy', 147 optimizer=gradient_descent.GradientDescentOptimizer(0.01)) 148 149 inputs = np.random.random((num_samples, timesteps, embedding_dim)) 150 initial_state = [ 151 np.random.random((num_samples, units)) for _ in range(num_states) 152 ] 153 targets = np.random.random((num_samples, units)) 154 model.train_on_batch([inputs] + initial_state, targets) 155 156 def DISABLED_test_specify_initial_state_non_keras_tensor(self): 157 num_states = 2 158 timesteps = 3 159 embedding_dim = 4 160 units = 3 161 num_samples = 2 162 163 # Test with non-Keras tensor 164 inputs = keras.Input((timesteps, embedding_dim)) 165 initial_state = [ 166 keras.backend.random_normal_variable((num_samples, units), 0, 1) 167 for _ in range(num_states) 168 ] 169 layer = rnn.LSTM(units) 170 output = layer(inputs, initial_state=initial_state) 171 172 model = keras.models.Model(inputs, output) 173 model.compile( 174 loss='categorical_crossentropy', 175 optimizer=gradient_descent.GradientDescentOptimizer(0.01)) 176 177 inputs = np.random.random((num_samples, timesteps, embedding_dim)) 178 targets = np.random.random((num_samples, units)) 179 model.train_on_batch(inputs, targets) 180 181 def test_reset_states_with_values(self): 182 num_states = 2 183 timesteps = 3 184 embedding_dim = 4 185 units = 3 186 num_samples = 2 187 188 layer = rnn.LSTM(units, stateful=True) 189 layer.build((num_samples, timesteps, embedding_dim)) 190 initial_weight_count = len(layer.weights) 191 layer.reset_states() 192 assert len(layer.states) == num_states 193 assert layer.states[0] is not None 194 self.assertAllClose( 195 keras.backend.eval(layer.states[0]), 196 np.zeros(keras.backend.int_shape(layer.states[0])), 197 atol=1e-4) 198 state_shapes = [keras.backend.int_shape(state) for state in layer.states] 199 values = [np.ones(shape) for shape in state_shapes] 200 if len(values) == 1: 201 values = values[0] 202 layer.reset_states(values) 203 self.assertAllClose( 204 keras.backend.eval(layer.states[0]), 205 np.ones(keras.backend.int_shape(layer.states[0])), 206 atol=1e-4) 207 208 # Test with invalid data 209 with self.assertRaises(ValueError): 210 layer.reset_states([1] * (len(layer.states) + 1)) 211 212 self.assertEqual(initial_weight_count, len(layer.weights)) 213 # Variables in "states" shouldn't show up in .weights 214 layer.states = nest.map_structure(variables.Variable, values) 215 layer.reset_states() 216 self.assertEqual(initial_weight_count, len(layer.weights)) 217 218 def test_specify_state_with_masking(self): 219 num_states = 2 220 timesteps = 3 221 embedding_dim = 4 222 units = 3 223 num_samples = 2 224 225 inputs = keras.Input((timesteps, embedding_dim)) 226 _ = keras.layers.Masking()(inputs) 227 initial_state = [keras.Input((units,)) for _ in range(num_states)] 228 output = rnn.LSTM(units)( 229 inputs, initial_state=initial_state) 230 231 model = keras.models.Model([inputs] + initial_state, output) 232 model.compile( 233 loss='categorical_crossentropy', 234 optimizer=gradient_descent.GradientDescentOptimizer(0.01)) 235 236 inputs = np.random.random((num_samples, timesteps, embedding_dim)) 237 initial_state = [ 238 np.random.random((num_samples, units)) for _ in range(num_states) 239 ] 240 targets = np.random.random((num_samples, units)) 241 model.train_on_batch([inputs] + initial_state, targets) 242 243 def test_return_state(self): 244 num_states = 2 245 timesteps = 3 246 embedding_dim = 4 247 units = 3 248 num_samples = 2 249 250 inputs = keras.Input(batch_shape=(num_samples, timesteps, embedding_dim)) 251 masked = keras.layers.Masking()(inputs) 252 layer = rnn.LSTM(units, return_state=True, stateful=True) 253 outputs = layer(masked) 254 state = outputs[1:] 255 assert len(state) == num_states 256 model = keras.models.Model(inputs, state[0]) 257 258 inputs = np.random.random((num_samples, timesteps, embedding_dim)) 259 state = model.predict(inputs) 260 self.assertAllClose(keras.backend.eval(layer.states[0]), state, atol=1e-4) 261 262 def test_state_reuse(self): 263 timesteps = 3 264 embedding_dim = 4 265 units = 3 266 num_samples = 2 267 268 inputs = keras.Input(batch_shape=(num_samples, timesteps, embedding_dim)) 269 layer = rnn.LSTM( 270 units, return_state=True, return_sequences=True) 271 outputs = layer(inputs) 272 output, state = outputs[0], outputs[1:] 273 output = rnn.LSTM(units)(output, initial_state=state) 274 model = keras.models.Model(inputs, output) 275 276 inputs = np.random.random((num_samples, timesteps, embedding_dim)) 277 model.predict(inputs) 278 279 def test_initial_states_as_other_inputs(self): 280 timesteps = 3 281 embedding_dim = 4 282 units = 3 283 num_samples = 2 284 num_states = 2 285 layer_class = rnn.LSTM 286 287 # Test with Keras tensor 288 main_inputs = keras.Input((timesteps, embedding_dim)) 289 initial_state = [keras.Input((units,)) for _ in range(num_states)] 290 inputs = [main_inputs] + initial_state 291 292 layer = layer_class(units) 293 output = layer(inputs) 294 assert initial_state[0] in layer._inbound_nodes[0].input_tensors 295 296 model = keras.models.Model(inputs, output) 297 model.compile( 298 loss='categorical_crossentropy', 299 optimizer=gradient_descent.GradientDescentOptimizer(0.01)) 300 301 main_inputs = np.random.random((num_samples, timesteps, embedding_dim)) 302 initial_state = [ 303 np.random.random((num_samples, units)) for _ in range(num_states) 304 ] 305 targets = np.random.random((num_samples, units)) 306 model.train_on_batch([main_inputs] + initial_state, targets) 307 308 def test_unified_lstm_feature_parity_with_canonical_lstm(self): 309 with context.eager_mode(): 310 # Run this test under eager only due to b/120160788 for model.set_weights. 311 input_shape = 10 312 rnn_state_size = 8 313 timestep = 4 314 batch = 20 315 316 (x_train, y_train), _ = testing_utils.get_test_data( 317 train_samples=batch, 318 test_samples=0, 319 input_shape=(timestep, input_shape), 320 num_classes=rnn_state_size) 321 y_train = keras.utils.to_categorical(y_train, rnn_state_size) 322 323 inputs = keras.layers.Input( 324 shape=[timestep, input_shape], dtype=dtypes.float32) 325 lstm_layer = rnn_v1.LSTM(rnn_state_size, 326 recurrent_activation='sigmoid') 327 output = lstm_layer(inputs) 328 lstm_model = keras.models.Model(inputs, output) 329 weights = lstm_model.get_weights() 330 y_1 = lstm_model.predict(x_train) 331 lstm_model.compile('rmsprop', 'mse') 332 lstm_model.fit(x_train, y_train) 333 y_2 = lstm_model.predict(x_train) 334 335 with test_util.device(use_gpu=True): 336 cudnn_layer = rnn.LSTM(rnn_state_size) 337 cudnn_model = keras.models.Model(inputs, cudnn_layer(inputs)) 338 cudnn_model.set_weights(weights) 339 y_3 = cudnn_model.predict(x_train) 340 cudnn_model.compile('rmsprop', 'mse') 341 cudnn_model.fit(x_train, y_train) 342 y_4 = cudnn_model.predict(x_train) 343 344 self.assertAllClose(y_1, y_3, rtol=1e-5, atol=2e-5) 345 self.assertAllClose(y_2, y_4, rtol=1e-5, atol=2e-5) 346 347 @parameterized.named_parameters(('v0', 0), ('v1', 1), ('v2', 2)) 348 def DISABLED_test_implementation_mode_LSTM(self, implementation_mode): 349 num_samples = 2 350 timesteps = 3 351 embedding_dim = 4 352 units = 2 353 testing_utils.layer_test( 354 rnn.LSTM, 355 kwargs={ 356 'units': units, 357 'implementation': implementation_mode 358 }, 359 input_shape=(num_samples, timesteps, embedding_dim)) 360 361 layer_class = rnn.LSTM 362 k_constraint = keras.constraints.max_norm(0.01) 363 r_constraint = keras.constraints.max_norm(0.01) 364 b_constraint = keras.constraints.max_norm(0.01) 365 layer = layer_class( 366 5, 367 return_sequences=False, 368 weights=None, 369 input_shape=(None, embedding_dim), 370 kernel_constraint=k_constraint, 371 recurrent_constraint=r_constraint, 372 bias_constraint=b_constraint) 373 layer.build((None, None, embedding_dim)) 374 self.assertEqual(layer.cell.kernel.constraint, k_constraint) 375 self.assertEqual(layer.cell.recurrent_kernel.constraint, r_constraint) 376 self.assertEqual(layer.cell.bias.constraint, b_constraint) 377 378 layer_class = rnn.LSTM 379 inputs = np.random.random((2, 3, 4)) 380 targets = np.abs(np.random.random((2, 3, 5))) 381 targets /= targets.sum(axis=-1, keepdims=True) 382 model = keras.models.Sequential() 383 model.add(keras.layers.Masking(input_shape=(3, 4))) 384 model.add(layer_class(units=5, return_sequences=True, unroll=False)) 385 model.compile( 386 loss='categorical_crossentropy', 387 optimizer=gradient_descent.GradientDescentOptimizer(0.01)) 388 model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1) 389 390 def test_masking_with_stacking_LSTM(self): 391 inputs = np.random.random((2, 3, 4)) 392 targets = np.abs(np.random.random((2, 3, 5))) 393 targets /= targets.sum(axis=-1, keepdims=True) 394 model = keras.models.Sequential() 395 model.add(keras.layers.Masking(input_shape=(3, 4))) 396 model.add(rnn.LSTM(10, return_sequences=True, unroll=False)) 397 model.add(rnn.LSTM(5, return_sequences=True, unroll=False)) 398 model.compile( 399 loss='categorical_crossentropy', 400 optimizer=gradient_descent.GradientDescentOptimizer(0.01)) 401 model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1) 402 403 @parameterized.named_parameters( 404 # test_name, time_major, go_backwards 405 ('normal', False, False), 406 ('time_major', True, False), 407 ('go_backwards', False, True), 408 ('both', True, True), 409 ) 410 def test_time_major_and_go_backward(self, time_major, go_backwards): 411 input_shape = 10 412 rnn_state_size = 8 413 timestep = 4 414 batch = 100 415 416 x_train = np.random.random((batch, timestep, input_shape)) 417 418 def build_model(layer_cls): 419 inputs = keras.layers.Input( 420 shape=[timestep, input_shape], dtype=dtypes.float32) 421 layer = layer_cls(rnn_state_size, 422 recurrent_activation='sigmoid', 423 time_major=time_major, 424 return_sequences=True, 425 go_backwards=go_backwards) 426 if time_major: 427 converted_input = keras.layers.Lambda( 428 lambda t: array_ops.transpose(t, [1, 0, 2]))(inputs) 429 outputs = layer(converted_input) 430 outputs = keras.layers.Lambda( 431 lambda t: array_ops.transpose(t, [1, 0, 2]))(outputs) 432 else: 433 outputs = layer(inputs) 434 return keras.models.Model(inputs, outputs) 435 436 lstm_model = build_model(rnn_v1.LSTM) 437 y_ref = lstm_model.predict(x_train) 438 weights = lstm_model.get_weights() 439 440 unified_lstm_model = build_model(rnn.LSTM) 441 unified_lstm_model.set_weights(weights) 442 y = unified_lstm_model.predict(x_train) 443 444 self.assertAllClose(y, y_ref) 445 446 input_shape = 10 447 rnn_state_size = 8 448 output_shape = 8 449 timestep = 4 450 batch = 100 451 epoch = 10 452 453 (x_train, y_train), _ = testing_utils.get_test_data( 454 train_samples=batch, 455 test_samples=0, 456 input_shape=(timestep, input_shape), 457 num_classes=output_shape) 458 y_train = keras.utils.to_categorical(y_train, output_shape) 459 460 layer = rnn.LSTM(rnn_state_size) 461 462 inputs = keras.layers.Input( 463 shape=[timestep, input_shape], dtype=dtypes.float32) 464 465 outputs = layer(inputs) 466 model = keras.models.Model(inputs, outputs) 467 model.compile('rmsprop', loss='mse') 468 model.fit(x_train, y_train, epochs=epoch) 469 model.evaluate(x_train, y_train) 470 model.predict(x_train) 471 472 @parameterized.named_parameters( 473 # test_name, use_bias, bias_initializer, activation 474 ('normal', True, 'zeros'), 475 ('no_bias', False, 'zeros'), 476 ('random_bias', True, 'random_uniform'), 477 ) 478 def test_unified_lstm_model_save_load(self, use_bias, bias_initializer): 479 temp_dir = self.get_temp_dir() 480 self.addCleanup(shutil.rmtree, temp_dir) 481 h5_path = os.path.join(temp_dir, 'test.h5') 482 483 batch = 10 484 timestep = 3 485 input_dim = 5 486 units = 2 487 488 x = np.random.random((batch, timestep, input_dim)) 489 490 def build_model(): 491 inputs = keras.layers.Input( 492 shape=[timestep, input_dim], dtype=dtypes.float32) 493 layer = rnn.LSTM( 494 units, 495 use_bias=use_bias, 496 bias_initializer=bias_initializer) 497 output = layer(inputs) 498 return keras.models.Model(inputs, output), layer 499 500 model, layer = build_model() 501 y_ref = model.predict(x) 502 model.save_weights(h5_path) 503 504 cloned_model, new_layer = build_model() 505 cloned_model.load_weights(h5_path) 506 y = cloned_model.predict(x) 507 508 self.assertAllClose(y, y_ref) 509 self.assertAllClose(layer.get_weights(), new_layer.get_weights()) 510 511 def test_unified_lstm_output_on_multiple_kernel(self): 512 input_shape = 10 513 rnn_state_size = 8 514 timestep = 4 515 batch = 100 516 517 x_train = np.random.random((batch, timestep, input_shape)) 518 519 inputs = keras.layers.Input( 520 shape=[timestep, input_shape], dtype=dtypes.float32) 521 with test_util.device(use_gpu=False): 522 layer = rnn.LSTM(rnn_state_size) 523 output = layer(inputs) 524 cpu_model = keras.models.Model(inputs, output) 525 weights = cpu_model.get_weights() 526 y_1 = cpu_model.predict(x_train) 527 528 with test_util.device(use_gpu=True): 529 layer = rnn.LSTM(rnn_state_size) 530 output = layer(inputs) 531 gpu_model = keras.models.Model(inputs, output) 532 gpu_model.set_weights(weights) 533 y_2 = gpu_model.predict(x_train) 534 535 # Note that CuDNN uses 'sigmoid' as activation, so the unified LSTM uses 536 # 'sigmoid' as default. Construct the canonical LSTM with sigmoid to achieve 537 # the same output. 538 with test_util.device(use_gpu=True): 539 layer = rnn_v1.LSTM(rnn_state_size, recurrent_activation='sigmoid') 540 output = layer(inputs) 541 canonical_model = keras.models.Model(inputs, output) 542 # Remove the extra cudnn bias since canonical lstm will not use it. 543 canonical_model.set_weights(weights[:3]) 544 y_3 = canonical_model.predict(x_train) 545 546 self.assertAllClose(y_1, y_2) 547 self.assertAllClose(y_2, y_3) 548 549 def DISABLED_test_return_sequences_LSTM(self): 550 num_samples = 2 551 timesteps = 3 552 embedding_dim = 4 553 units = 2 554 testing_utils.layer_test( 555 rnn.LSTM, 556 kwargs={ 557 'units': units, 558 'return_sequences': True 559 }, 560 input_shape=(num_samples, timesteps, embedding_dim)) 561 562 def test_regularizers_LSTM(self): 563 embedding_dim = 4 564 layer_class = rnn.LSTM 565 layer = layer_class( 566 5, 567 return_sequences=False, 568 weights=None, 569 input_shape=(None, embedding_dim), 570 kernel_regularizer=keras.regularizers.l1(0.01), 571 recurrent_regularizer=keras.regularizers.l1(0.01), 572 bias_regularizer='l2', 573 activity_regularizer='l1') 574 layer.build((None, None, 2)) 575 self.assertEqual(len(layer.losses), 3) 576 x = keras.backend.variable(np.ones((2, 3, 2))) 577 layer(x) 578 if context.executing_eagerly(): 579 self.assertEqual(len(layer.losses), 4) 580 else: 581 self.assertEqual(len(layer.get_losses_for(x)), 1) 582 583 def test_statefulness_LSTM(self): 584 num_samples = 2 585 timesteps = 3 586 embedding_dim = 4 587 units = 2 588 layer_class = rnn.LSTM 589 model = keras.models.Sequential() 590 model.add( 591 keras.layers.Embedding( 592 4, 593 embedding_dim, 594 mask_zero=True, 595 input_length=timesteps, 596 batch_input_shape=(num_samples, timesteps))) 597 layer = layer_class( 598 units, return_sequences=False, stateful=True, weights=None) 599 model.add(layer) 600 model.compile(optimizer=gradient_descent.GradientDescentOptimizer(0.01), 601 loss='mse', run_eagerly=testing_utils.should_run_eagerly()) 602 out1 = model.predict(np.ones((num_samples, timesteps))) 603 self.assertEqual(out1.shape, (num_samples, units)) 604 605 # train once so that the states change 606 model.train_on_batch( 607 np.ones((num_samples, timesteps)), np.ones((num_samples, units))) 608 out2 = model.predict(np.ones((num_samples, timesteps))) 609 610 # if the state is not reset, output should be different 611 self.assertNotEqual(out1.max(), out2.max()) 612 613 # check that output changes after states are reset 614 # (even though the model itself didn't change) 615 layer.reset_states() 616 out3 = model.predict(np.ones((num_samples, timesteps))) 617 self.assertNotEqual(out2.max(), out3.max()) 618 619 # check that container-level reset_states() works 620 model.reset_states() 621 out4 = model.predict(np.ones((num_samples, timesteps))) 622 self.assertAllClose(out3, out4, atol=1e-5) 623 624 # check that the call to `predict` updated the states 625 out5 = model.predict(np.ones((num_samples, timesteps))) 626 self.assertNotEqual(out4.max(), out5.max()) 627 628 # Check masking 629 layer.reset_states() 630 631 left_padded_input = np.ones((num_samples, timesteps)) 632 left_padded_input[0, :1] = 0 633 left_padded_input[1, :2] = 0 634 out6 = model.predict(left_padded_input) 635 636 layer.reset_states() 637 638 right_padded_input = np.ones((num_samples, timesteps)) 639 right_padded_input[0, -1:] = 0 640 right_padded_input[1, -2:] = 0 641 out7 = model.predict(right_padded_input) 642 643 self.assertAllClose(out7, out6, atol=1e-5) 644 645 def test_stateful_LSTM_training(self): 646 # See b/123587692 for more context. 647 vocab_size = 20 648 embedding_dim = 10 649 batch_size = 8 650 timestep = 12 651 units = 5 652 x = np.random.randint(0, vocab_size, size=(batch_size, timestep)) 653 y = np.random.randint(0, vocab_size, size=(batch_size, timestep)) 654 655 model = keras.Sequential([ 656 keras.layers.Embedding(vocab_size, embedding_dim, 657 batch_input_shape=[batch_size, timestep]), 658 rnn.LSTM(units, return_sequences=True, stateful=True), 659 keras.layers.Dense(vocab_size) 660 ]) 661 model.compile(optimizer='adam', 662 loss='sparse_categorical_crossentropy', 663 run_eagerly=testing_utils.should_run_eagerly()) 664 model.fit(x, y, epochs=1, shuffle=False) 665 666 def test_dropout_LSTM(self): 667 num_samples = 2 668 timesteps = 3 669 embedding_dim = 4 670 units = 2 671 testing_utils.layer_test( 672 rnn.LSTM, 673 kwargs={ 674 'units': units, 675 'dropout': 0.1, 676 'recurrent_dropout': 0.1 677 }, 678 input_shape=(num_samples, timesteps, embedding_dim)) 679 680 681 class LSTMLayerGraphOnlyTest(test.TestCase): 682 683 # Need session for test 684 @test_util.run_deprecated_v1 685 def test_unifiedLSTM(self): 686 input_shape = 10 687 rnn_state_size = 8 688 output_shape = 8 689 timestep = 4 690 batch = 100 691 epoch = 1 692 693 with self.cached_session(config=_config, use_gpu=True) as sess: 694 (x_train, y_train), _ = testing_utils.get_test_data( 695 train_samples=batch, 696 test_samples=0, 697 input_shape=(timestep, input_shape), 698 num_classes=output_shape) 699 y_train = keras.utils.to_categorical(y_train, output_shape) 700 701 layer = rnn.LSTM(rnn_state_size, return_runtime=True) 702 703 inputs = array_ops.placeholder( 704 dtypes.float32, shape=(None, timestep, input_shape), name='inputs') 705 predict = array_ops.placeholder( 706 dtypes.float32, shape=(None, output_shape), name='predict') 707 708 outputs, runtime = layer(inputs) 709 loss = losses.softmax_cross_entropy(predict, outputs) 710 optimizer = gradient_descent.GradientDescentOptimizer(0.001) 711 train_op = optimizer.minimize(loss) 712 713 sess.run([variables.global_variables_initializer()]) 714 existing_loss = 0 715 for _ in range(epoch): 716 loss_value, _, runtime_value = sess.run([loss, train_op, runtime], { 717 inputs: x_train, 718 predict: y_train 719 }) 720 if test.is_gpu_available(): 721 self.assertEqual(runtime_value, b'cudnn') 722 else: 723 self.assertEqual(runtime_value, b'cpu') 724 # Make sure the loss is updated for every epoch 725 # (layer weights properly updated). 726 self.assertNotEqual(existing_loss, loss_value) 727 existing_loss = loss_value 728 729 # Need session for test 730 @test_util.run_deprecated_v1 731 def test_unifiedLSTM_with_cond(self): 732 # This test is to demonstrate the graph rewrite of grappler plugin under 733 # the condition that the function returns different number of internal 734 # states. 735 input_shape = 10 736 rnn_state_size = 8 737 output_shape = 8 738 timestep = 4 739 batch = 100 740 epoch = 1 741 742 with self.cached_session(config=_config, use_gpu=True) as sess: 743 (x_train, y_train), _ = testing_utils.get_test_data( 744 train_samples=batch, 745 test_samples=0, 746 input_shape=(timestep, input_shape), 747 num_classes=output_shape) 748 y_train = keras.utils.to_categorical(y_train, output_shape) 749 750 layer = rnn.LSTM(rnn_state_size, return_runtime=True) 751 752 inputs = array_ops.placeholder( 753 dtypes.float32, shape=(None, timestep, input_shape), name='inputs') 754 predict = array_ops.placeholder( 755 dtypes.float32, shape=(None, output_shape), name='predict') 756 757 zeros = array_ops.zeros([batch, output_shape]) 758 dummy_runtime = constant_op.constant( 759 'unknown', dtype=dtypes.string, name='runtime') 760 a = constant_op.constant(0) 761 b = constant_op.constant(1) 762 # Will always run the lstm layer. 763 outputs, runtime = control_flow_ops.cond( 764 gen_math_ops.less(a, b), 765 lambda: layer(inputs), 766 lambda: (zeros, dummy_runtime)) 767 loss = losses.softmax_cross_entropy(predict, outputs) 768 optimizer = gradient_descent.GradientDescentOptimizer(0.001) 769 train_op = optimizer.minimize(loss) 770 771 sess.run([variables.global_variables_initializer()]) 772 existing_loss = 0 773 774 for _ in range(epoch): 775 loss_value, _, runtime_value = sess.run([loss, train_op, runtime], { 776 inputs: x_train, 777 predict: y_train 778 }) 779 if test.is_gpu_available(): 780 self.assertEqual(runtime_value, b'cudnn') 781 else: 782 self.assertEqual(runtime_value, b'cpu') 783 # Make sure the loss is updated for every epoch 784 # (layer weights properly updated). 785 self.assertNotEqual(existing_loss, loss_value) 786 existing_loss = loss_value 787 788 789 class UnifiedLSTMPerformanceTest(test.Benchmark): 790 791 def _measure_performance(self, test_config, model, x_train, y_train): 792 batch = test_config['batch'] 793 epoch = test_config['epoch'] 794 warmup_epoch = test_config['warmup_epoch'] 795 796 # warm up the model 797 model.fit(x_train, y_train, batch_size=batch, epochs=warmup_epoch) 798 start_time = time.time() 799 model.fit(x_train, y_train, batch_size=batch, epochs=epoch - warmup_epoch) 800 end_time = time.time() 801 return (end_time - start_time) / (epoch - warmup_epoch) 802 803 def _time_performance_run_cudnn_lstm(self, test_config, x_train, y_train): 804 # Get the performance number for standard Cudnn LSTM 805 input_shape = test_config['input_shape'] 806 rnn_state_size = test_config['rnn_state_size'] 807 timestep = test_config['timestep'] 808 809 cudnn_lstm_layer = keras.layers.CuDNNLSTM(rnn_state_size) 810 inputs = keras.layers.Input( 811 shape=[timestep, input_shape], dtype=dtypes.float32) 812 813 outputs = cudnn_lstm_layer(inputs) 814 model = keras.models.Model(inputs, outputs) 815 model.compile('sgd', 'mse') 816 817 sec_per_epoch = self._measure_performance( 818 test_config, model, x_train, y_train) 819 logging.info('Average performance for %s per epoch is: %s', 820 'CuDNN LSTM', sec_per_epoch) 821 return sec_per_epoch 822 823 def _time_performance_run_unifed_lstm_gpu( 824 self, test_config, x_train, y_train): 825 # Get performance number for Unified_LSTM with grappler swap the impl 826 input_shape = test_config['input_shape'] 827 rnn_state_size = test_config['rnn_state_size'] 828 timestep = test_config['timestep'] 829 830 layer = rnn.LSTM(rnn_state_size) 831 inputs = keras.layers.Input( 832 shape=[timestep, input_shape], dtype=dtypes.float32) 833 834 outputs = layer(inputs) 835 model = keras.models.Model(inputs, outputs) 836 model.compile('sgd', 'mse') 837 838 sec_per_epoch = self._measure_performance( 839 test_config, model, x_train, y_train) 840 logging.info('Average performance for %s per epoch is: %s', 841 'Unified LSTM', sec_per_epoch) 842 return sec_per_epoch 843 844 def _time_performance_run_normal_lstm( 845 self, test_config, x_train, y_train): 846 # Get performance number for standard LSTM on GPU. 847 input_shape = test_config['input_shape'] 848 rnn_state_size = test_config['rnn_state_size'] 849 timestep = test_config['timestep'] 850 851 layer = rnn_v1.LSTM(rnn_state_size) 852 inputs = keras.layers.Input( 853 shape=[timestep, input_shape], dtype=dtypes.float32) 854 855 outputs = layer(inputs) 856 model = keras.models.Model(inputs, outputs) 857 model.compile('sgd', 'mse') 858 859 sec_per_epoch = self._measure_performance( 860 test_config, model, x_train, y_train) 861 logging.info('Average performance for %s per epoch is: %s', 862 'Normal LSTM', sec_per_epoch) 863 return sec_per_epoch 864 865 def _benchmark_performance_with_standard_cudnn_impl(self): 866 if not test.is_gpu_available(): 867 self.skipTest('performance test will only run on GPU') 868 869 mode = 'eager' if context.executing_eagerly() else 'graph' 870 batch = 64 871 num_batch = 10 872 test_config = { 873 'input_shape': 128, 874 'rnn_state_size': 64, 875 'output_shape': 64, 876 'timestep': 50, 877 'batch': batch, 878 'epoch': 20, 879 # The performance for warmup epoch is ignored. 880 'warmup_epoch': 1, 881 } 882 (x_train, y_train), _ = testing_utils.get_test_data( 883 train_samples=(batch * num_batch), 884 test_samples=0, 885 input_shape=(test_config['timestep'], test_config['input_shape']), 886 num_classes=test_config['output_shape']) 887 y_train = keras.utils.to_categorical(y_train, test_config['output_shape']) 888 889 cudnn_sec_per_epoch = self._time_performance_run_cudnn_lstm( 890 test_config, x_train, y_train) 891 unified_lstm_sec_per_epoch = self._time_performance_run_unifed_lstm_gpu( 892 test_config, x_train, y_train) 893 normal_lstm_sec_per_epoch = self._time_performance_run_normal_lstm( 894 test_config, x_train, y_train) 895 896 cudnn_vs_unified = cudnn_sec_per_epoch / unified_lstm_sec_per_epoch 897 unified_vs_normal = normal_lstm_sec_per_epoch / unified_lstm_sec_per_epoch 898 899 self.report_benchmark(name='keras_cudnn_lstm_' + mode, 900 wall_time=cudnn_sec_per_epoch, 901 iters=test_config['epoch'], 902 extras=test_config) 903 self.report_benchmark(name='keras_unified_lstm_' + mode, 904 wall_time=unified_lstm_sec_per_epoch, 905 iters=test_config['epoch'], 906 extras=test_config) 907 self.report_benchmark(name='keras_canonical_lstm_' + mode, 908 wall_time=normal_lstm_sec_per_epoch, 909 iters=test_config['epoch'], 910 extras=test_config) 911 912 logging.info('Expect the performance of Unified LSTM is within 80% of ' 913 'CuDNN LSTM, got {0:.2f}%'.format(cudnn_vs_unified * 100)) 914 logging.info('Expect the performance of Unified LSTM is more than 5 times' 915 ' of normal LSTM, got {0:.2f}'.format(unified_vs_normal)) 916 917 def benchmark_performance_graph(self): 918 with context.graph_mode(), session_lib.Session(config=_config): 919 self._benchmark_performance_with_standard_cudnn_impl() 920 921 def benchmark_performance_eager(self): 922 with context.eager_mode(): 923 self._benchmark_performance_with_standard_cudnn_impl() 924 925 926 if __name__ == '__main__': 927 test.main() 928
