def test_deepcopy(self): if not context.executing_eagerly(): self.skipTest('v2-only test') original_layer = rnn.GRU(5) copied_layer = copy.deepcopy(original_layer) self.assertEqual(copied_layer.units, 5) self.assertEqual(original_layer.get_config(), original_layer.get_config()) # Copy layer before layer call on inputs without weight initialization. inputs = np.random.normal(size=[32, 10, 8]).astype(np.float32) original_layer = rnn.GRU(4) copied_layer = copy.deepcopy(original_layer) outputs = original_layer(inputs) copied_outputs = copied_layer(inputs) self.assertNotAllClose(self.evaluate(outputs), self.evaluate(copied_outputs)) # Copy layer after layer call on inputs with weight initialization. original_layer = rnn.GRU(4) outputs = original_layer(inputs) copied_layer = copy.deepcopy(original_layer) copied_outputs = copied_layer(inputs) self.assertAllClose(self.evaluate(outputs), self.evaluate(copied_outputs))
def test_stacking_GRU(self): inputs = np.random.random((2, 3, 4)) targets = np.abs(np.random.random((2, 3, 5))) targets /= targets.sum(axis=-1, keepdims=True) model = keras.models.Sequential() model.add(rnn.GRU(10, return_sequences=True, unroll=False)) model.add(rnn.GRU(5, return_sequences=True, unroll=False)) model.compile( loss='categorical_crossentropy', optimizer=gradient_descent.GradientDescentOptimizer(0.01)) model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)
def test_explicit_device_with_go_backward_and_mask(self): if test.is_built_with_rocm(): self.skipTest('Skipping the test as ROCm MIOpen does not ' 'support padded input yet.') batch_size = 8 timestep = 7 masksteps = 5 units = 4 inputs = np.random.randn(batch_size, timestep, units).astype(np.float32) mask = np.ones((batch_size, timestep)).astype(np.bool) mask[:, masksteps:] = 0 # Test for V1 behavior. lstm_v1 = rnn_v1.GRU(units, return_sequences=True, go_backwards=True) with test_util.device(use_gpu=True): outputs_masked_v1 = lstm_v1(inputs, mask=constant_op.constant(mask)) outputs_trimmed_v1 = lstm_v1(inputs[:, :masksteps]) self.assertAllClose(outputs_masked_v1[:, -masksteps:], outputs_trimmed_v1) # Test for V2 behavior. lstm = rnn.GRU(units, return_sequences=True, go_backwards=True) with test_util.device(use_gpu=True): outputs_masked = lstm(inputs, mask=constant_op.constant(mask)) outputs_trimmed = lstm(inputs[:, :masksteps]) self.assertAllClose(outputs_masked[:, -masksteps:], outputs_trimmed)
def test_keras_model_with_gru(self): input_shape = 10 rnn_state_size = 8 output_shape = 8 timestep = 4 batch = 100 epoch = 10 (x_train, y_train), _ = testing_utils.get_test_data( train_samples=batch, test_samples=0, input_shape=(timestep, input_shape), num_classes=output_shape) y_train = keras.utils.to_categorical(y_train, output_shape) layer = rnn.GRU(rnn_state_size) inputs = keras.layers.Input( shape=[timestep, input_shape], dtype=dtypes.float32) outputs = layer(inputs) model = keras.models.Model(inputs, outputs) model.compile('rmsprop', loss='mse') model.fit(x_train, y_train, epochs=epoch) model.evaluate(x_train, y_train) model.predict(x_train)
def test_GRU_runtime_with_cond(self): # This test is to demonstrate the graph rewrite of grappler plugin under # the condition that the function returns different number of internal # states. layer = rnn.GRU(self.rnn_state_size, return_runtime=True) inputs = keras.layers.Input(shape=[self.timestep, self.input_shape], dtype=dtypes.float32) zeros = array_ops.zeros([self.batch, self.output_shape]) dummy_runtime = rnn._runtime(rnn._RUNTIME_UNKNOWN) a = constant_op.constant(0) b = constant_op.constant(1) # Will always run the GRU layer. outputs, runtime = control_flow_ops.cond( gen_math_ops.less(a, b), lambda: layer(inputs), lambda: (zeros, dummy_runtime)) # Expand the runtime so that it is a 1D tensor instead of scalar. # TF model does not work with scalar model output, specially during # aggregation. runtime = keras.layers.Lambda( lambda x: array_ops.expand_dims(x, axis=-1))(runtime) model = keras.models.Model(inputs=inputs, outputs=[outputs, runtime]) self._test_runtime_with_model(model)
def test_masking_with_stacking_GRU(self): if test.is_built_with_rocm(): self.skipTest('Skipping the test as ROCm MIOpen does not ' 'support padded input yet.') inputs = np.random.random((2, 3, 4)) targets = np.abs(np.random.random((2, 3, 5))) targets /= targets.sum(axis=-1, keepdims=True) model = keras.models.Sequential() model.add(keras.layers.Masking(input_shape=(3, 4))) model.add(rnn.GRU(10, return_sequences=True, unroll=False)) model.add(rnn.GRU(5, return_sequences=True, unroll=False)) model.compile( loss='categorical_crossentropy', optimizer=gradient_descent.GradientDescentOptimizer(0.01)) model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)
def test_with_fully_masked_inputs(self): num_samples = 8 timestep = 5 embedding_dim = 4 vocab_size = 20 units = 2 inputs = np.random.randint(0, vocab_size, size=(num_samples, timestep)) # Set the first inputs to be fully zero. inputs[0, :] = 0.0 model = keras.models.Sequential() model.add( keras.layers.Embedding( vocab_size, embedding_dim, mask_zero=True, input_length=timestep, batch_input_shape=(num_samples, timestep))) layer = rnn.GRU(units) model.add(layer) model.compile( optimizer=gradient_descent.GradientDescentOptimizer(0.01), loss='mse', run_eagerly=testing_utils.should_run_eagerly()) # Make sure it doesn't crash with cudnn kernel. model.predict(inputs)
def test_explicit_device_with_go_backward_and_mask(self): batch_size = 8 timestep = 7 masksteps = 5 units = 4 inputs = np.random.randn(batch_size, timestep, units).astype(np.float32) mask = np.ones((batch_size, timestep)).astype(np.bool) mask[:, masksteps:] = 0 # Test for V1 behavior. lstm_v1 = rnn_v1.GRU(units, return_sequences=True, go_backwards=True) with test_util.device(use_gpu=True): outputs_masked_v1 = lstm_v1(inputs, mask=constant_op.constant(mask)) outputs_trimmed_v1 = lstm_v1(inputs[:, :masksteps]) self.assertAllClose(outputs_masked_v1[:, -masksteps:], outputs_trimmed_v1) # Test for V2 behavior. lstm = rnn.GRU(units, return_sequences=True, go_backwards=True) with test_util.device(use_gpu=True): outputs_masked = lstm(inputs, mask=constant_op.constant(mask)) outputs_trimmed = lstm(inputs[:, :masksteps]) self.assertAllClose(outputs_masked[:, -masksteps:], outputs_trimmed)
def build_model(): inputs = keras.layers.Input(shape=[timestep, input_dim], dtype=dtypes.float32) layer = rnn.GRU(units, use_bias=use_bias, bias_initializer=bias_initializer) output = layer(inputs) return keras.models.Model(inputs, output), layer
def test_GRU_runtime_with_mask(self): if test.is_built_with_rocm(): self.skipTest('Skipping the test as ROCm MIOpen does not ' 'support padded input yet.') # Masking will affect which backend is selected based on whether the mask # is strictly right padded. layer = rnn.GRU(self.rnn_state_size, return_runtime=True) inputs = keras.layers.Input(shape=[self.timestep, self.input_shape], dtype=dtypes.float32) masked_inputs = keras.layers.Masking()(inputs) outputs, runtime = layer(masked_inputs) # Expand the runtime so that it is a 1D tensor instead of scalar. # TF model does not work with scalar model output, specially during # aggregation. runtime = keras.layers.Lambda( lambda x: array_ops.expand_dims(x, axis=-1))(runtime) model = keras.models.Model(inputs=inputs, outputs=[outputs, runtime]) (x_train, y_train), _ = testing_utils.get_test_data( train_samples=self.batch, test_samples=0, input_shape=(self.timestep, self.input_shape), num_classes=self.output_shape) y_train = np_utils.to_categorical(y_train, self.output_shape) model.compile(optimizer='sgd', loss=['categorical_crossentropy', None], run_eagerly=testing_utils.should_run_eagerly(), experimental_run_tf_function=testing_utils. should_run_tf_function()) model.fit(x_train, y_train) # Verify unpadded data. _, runtime_value = model.predict(x_train) if test.is_gpu_available(): self.assertEqual(runtime_value[0], rnn._RUNTIME_GPU) else: self.assertEqual(runtime_value[0], rnn._RUNTIME_CPU) # Update x/y to be right padded by setting the last timestep to 0 x_train[:, -1, :] = 0 y_train[:, -1] = 0 _, runtime_value = model.predict(x_train) if test.is_gpu_available(): self.assertEqual(runtime_value[0], rnn._RUNTIME_GPU) else: self.assertEqual(runtime_value[0], rnn._RUNTIME_CPU) # Further update x/y to be mix padded (masks in the middle), and verify # only cpu kernel can be selected. x_train[:, -3, :] = 0 y_train[:, -3] = 0 _, runtime_value = model.predict(x_train) self.assertEqual(runtime_value[0], rnn._RUNTIME_CPU)
def test_UnifiedGRU_with_cond(self): # This test is to demonstrate the graph rewrite of grappler plugin under # the condition that the function returns different number of internal # states. input_shape = 10 rnn_state_size = 8 output_shape = 8 timestep = 4 batch = 100 epoch = 1 with self.cached_session(config=_config, use_gpu=True) as sess: (x_train, y_train), _ = testing_utils.get_test_data( train_samples=batch, test_samples=0, input_shape=(timestep, input_shape), num_classes=output_shape) y_train = keras.utils.to_categorical(y_train, output_shape) layer = rnn.GRU(rnn_state_size, return_runtime=True) inputs = array_ops.placeholder(dtypes.float32, shape=(None, timestep, input_shape), name='inputs') predict = array_ops.placeholder(dtypes.float32, shape=(None, output_shape), name='predict') zeros = array_ops.zeros([batch, output_shape]) dummy_runtime = rnn._runtime(rnn._RUNTIME_UNKNOWN) a = constant_op.constant(0) b = constant_op.constant(1) # Will always run the GRU layer. outputs, runtime = control_flow_ops.cond( gen_math_ops.less(a, b), lambda: layer(inputs), lambda: (zeros, dummy_runtime)) loss = losses.softmax_cross_entropy(predict, outputs) optimizer = gradient_descent.GradientDescentOptimizer(0.001) train_op = optimizer.minimize(loss) sess.run([variables.global_variables_initializer()]) existing_loss = 0 for _ in range(epoch): loss_value, _, runtime_value = sess.run( [loss, train_op, runtime], { inputs: x_train, predict: y_train }) if test.is_gpu_available(): self.assertEqual(runtime_value, rnn._RUNTIME_GPU) else: self.assertEqual(runtime_value, rnn._RUNTIME_CPU) # Make sure the loss is updated for every epoch # (layer weights properly updated). self.assertNotEqual(existing_loss, loss_value) existing_loss = loss_value
def test_could_use_defun_backend(self, activation, recurrent_activation, recurrent_dropout, unroll, use_bias, reset_after): layer = rnn.GRU(1, activation=activation, recurrent_activation=recurrent_activation, recurrent_dropout=recurrent_dropout, unroll=unroll, use_bias=use_bias, reset_after=reset_after) self.assertFalse(layer.could_use_cudnn)
def test_dynamic_behavior_GRU(self): num_samples = 2 timesteps = 3 embedding_dim = 4 units = 2 layer = rnn.GRU(units, input_shape=(None, embedding_dim)) model = keras.models.Sequential() model.add(layer) model.compile(gradient_descent.GradientDescentOptimizer(0.001), 'mse') x = np.random.random((num_samples, timesteps, embedding_dim)) y = np.random.random((num_samples, units)) model.train_on_batch(x, y)
def test_unified_gru_output_on_multiple_kernel(self): input_shape = 10 rnn_state_size = 8 timestep = 4 batch = 100 x_train = np.random.random((batch, timestep, input_shape)) inputs = keras.layers.Input( shape=[timestep, input_shape], dtype=dtypes.float32) with test_util.device(use_gpu=False): layer = rnn.GRU(rnn_state_size) output = layer(inputs) cpu_model = keras.models.Model(inputs, output) weights = cpu_model.get_weights() y_1 = cpu_model.predict(x_train) with test_util.device(use_gpu=True): layer = rnn.GRU(rnn_state_size) output = layer(inputs) gpu_model = keras.models.Model(inputs, output) gpu_model.set_weights(weights) y_2 = gpu_model.predict(x_train) # Note that CuDNN uses 'sigmoid' as activation, so the unified GRU uses # 'sigmoid' as default. Construct the canonical GRU with sigmoid to achieve # the same output. with test_util.device(use_gpu=True): layer = rnn_v1.GRU(rnn_state_size, recurrent_activation='sigmoid', reset_after=True) output = layer(inputs) canonical_model = keras.models.Model(inputs, output) canonical_model.set_weights(weights) y_3 = canonical_model.predict(x_train) self.assertAllClose(y_1, y_2) self.assertAllClose(y_2, y_3)
def test_GRU_runtime(self): layer = rnn.GRU(self.rnn_state_size, return_runtime=True) inputs = keras.layers.Input(shape=[self.timestep, self.input_shape], dtype=dtypes.float32) outputs, runtime = layer(inputs) # Expand the runtime so that it is a 1D tensor instead of scalar. # TF model does not work with scalar model output, specially during # aggregation. runtime = keras.layers.Lambda( lambda x: array_ops.expand_dims(x, axis=-1))(runtime) model = keras.models.Model(inputs=inputs, outputs=[outputs, runtime]) self._test_runtime_with_model(model)
def test_gru_v2_feature_parity_with_canonical_gru(self): if test.is_built_with_rocm(): self.skipTest('Skipping the test as ROCm MIOpen does not ' 'support padded input yet.') input_shape = 10 rnn_state_size = 8 timestep = 4 batch = 20 (x_train, y_train), _ = testing_utils.get_test_data( train_samples=batch, test_samples=0, input_shape=(timestep, input_shape), num_classes=rnn_state_size, random_seed=random_seed.DEFAULT_GRAPH_SEED) y_train = np_utils.to_categorical(y_train, rnn_state_size) # For the last batch item of the test data, we filter out the last # timestep to simulate the variable length sequence and masking test. x_train[-2:, -1, :] = 0.0 y_train[-2:] = 0 inputs = keras.layers.Input(shape=[timestep, input_shape], dtype=dtypes.float32) masked_input = keras.layers.Masking()(inputs) gru_layer = rnn_v1.GRU(rnn_state_size, recurrent_activation='sigmoid', reset_after=True) output = gru_layer(masked_input) gru_model = keras.models.Model(inputs, output) weights = gru_model.get_weights() y_1 = gru_model.predict(x_train) gru_model.compile('rmsprop', 'mse') gru_model.fit(x_train, y_train) y_2 = gru_model.predict(x_train) with test_util.device(use_gpu=True): cudnn_layer = rnn.GRU(rnn_state_size, recurrent_activation='sigmoid', reset_after=True) cudnn_model = keras.models.Model(inputs, cudnn_layer(masked_input)) cudnn_model.set_weights(weights) y_3 = cudnn_model.predict(x_train) cudnn_model.compile('rmsprop', 'mse') cudnn_model.fit(x_train, y_train) y_4 = cudnn_model.predict(x_train) self.assertAllClose(y_1, y_3, rtol=2e-5, atol=2e-5) self.assertAllClose(y_2, y_4, rtol=2e-5, atol=2e-5)
def test_v1_session_behavior(self): # See b/139132348 for more details. x = np.random.uniform(size=(100, 4, 8)) y = np.random.uniform(size=(100, 1)) dataset = dataset_ops.Dataset.from_tensor_slices( (x, y)).shuffle(100).batch(32) inp = keras.layers.Input(shape=(4, 8)) layer = rnn.GRU(1)(inp) layer = keras.layers.Dense(1)(layer) model = keras.models.Model(inp, layer) model.compile(loss='mse', optimizer='sgd') model.fit(dataset)
def test_unified_gru_feature_parity_with_canonical_gru(self): with context.eager_mode(): # Run this test under eager only due to b/120160788 for model.set_weights. input_shape = 10 rnn_state_size = 8 timestep = 4 batch = 20 (x_train, y_train), _ = testing_utils.get_test_data( train_samples=batch, test_samples=0, input_shape=(timestep, input_shape), num_classes=rnn_state_size) y_train = keras.utils.to_categorical(y_train, rnn_state_size) # For the last batch item of the test data, we filter out the last # timestep to simulate the variable length sequence and masking test. x_train[-2:, -1, :] = 0.0 y_train[-2:] = 0 inputs = keras.layers.Input(shape=[timestep, input_shape], dtype=dtypes.float32) masked_input = keras.layers.Masking()(inputs) gru_layer = rnn_v1.GRU(rnn_state_size, recurrent_activation='sigmoid', reset_after=True) output = gru_layer(masked_input) gru_model = keras.models.Model(inputs, output) weights = gru_model.get_weights() y_1 = gru_model.predict(x_train) gru_model.compile('rmsprop', 'mse') gru_model.fit(x_train, y_train) y_2 = gru_model.predict(x_train) with test_util.device(use_gpu=True): cudnn_layer = rnn.GRU(rnn_state_size, recurrent_activation='sigmoid', reset_after=True) cudnn_model = keras.models.Model(inputs, cudnn_layer(masked_input)) cudnn_model.set_weights(weights) y_3 = cudnn_model.predict(x_train) cudnn_model.compile('rmsprop', 'mse') cudnn_model.fit(x_train, y_train) y_4 = cudnn_model.predict(x_train) self.assertAllClose(y_1, y_3, rtol=2e-5, atol=2e-5) self.assertAllClose(y_2, y_4, rtol=2e-5, atol=2e-5)
def test_unifiedGRU(self): input_shape = 10 rnn_state_size = 8 output_shape = 8 timestep = 4 batch = 100 epoch = 1 with self.cached_session(config=_config, use_gpu=True) as sess: (x_train, y_train), _ = testing_utils.get_test_data( train_samples=batch, test_samples=0, input_shape=(timestep, input_shape), num_classes=output_shape) y_train = keras.utils.to_categorical(y_train, output_shape) layer = rnn.GRU(rnn_state_size, return_runtime=True) inputs = array_ops.placeholder(dtypes.float32, shape=(None, timestep, input_shape), name='inputs') predict = array_ops.placeholder(dtypes.float32, shape=(None, output_shape), name='predict') outputs, runtime = layer(inputs) loss = losses.softmax_cross_entropy(predict, outputs) optimizer = gradient_descent.GradientDescentOptimizer(0.001) train_op = optimizer.minimize(loss) sess.run([variables.global_variables_initializer()]) existing_loss = 0 for _ in range(epoch): loss_value, _, runtime_value = sess.run( [loss, train_op, runtime], { inputs: x_train, predict: y_train }) if test.is_gpu_available(): self.assertEqual(runtime_value, rnn._RUNTIME_GPU) else: self.assertEqual(runtime_value, rnn._RUNTIME_CPU) # Make sure the loss is updated for every epoch # (layer weights properly updated). self.assertNotEqual(existing_loss, loss_value) existing_loss = loss_value
def test_stateful_GRU_training(self): # See b/123587692 for more context. vocab_size = 20 embedding_dim = 10 batch_size = 8 timestep = 12 units = 5 x = np.random.randint(0, vocab_size, size=(batch_size, timestep)) y = np.random.randint(0, vocab_size, size=(batch_size, timestep)) model = keras.Sequential([ keras.layers.Embedding(vocab_size, embedding_dim, batch_input_shape=[batch_size, timestep]), rnn.GRU(units, return_sequences=True, stateful=True), keras.layers.Dense(vocab_size) ]) model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', run_eagerly=testing_utils.should_run_eagerly()) model.fit(x, y, epochs=1, shuffle=False)
def test_in_tape(self): with self.test_session(config=_config): time_steps = 10 embedding_size = 11 gru_unit_size = 12 gru = rnn.GRU(gru_unit_size, return_sequences=True, return_state=True, recurrent_activation='sigmoid', recurrent_initializer='glorot_uniform') x = random_ops.random_uniform([1, time_steps, embedding_size]) y = random_ops.random_uniform([1, gru_unit_size]) with backprop.GradientTape() as tape: hidden_state = array_ops.zeros([1, gru_unit_size], dtype=dtypes.float32) _, state = gru(x, initial_state=hidden_state) loss = math_ops.reduce_mean(math_ops.square(state - y)) tape.gradient(loss, gru.variables)
def test_gru(self): self._test_layer(recurrent_v2.GRU(units=4, return_sequences=True), input_shape=(4, 4, 4))
def test_use_on_default_activation_with_gpu_kernel(self): layer = rnn.GRU(1, activation=nn.tanh) self.assertTrue(layer._could_use_gpu_kernel) layer = rnn.GRU(1, recurrent_activation=nn.sigmoid) self.assertTrue(layer._could_use_gpu_kernel)
class LayerCorrectnessTest(keras_parameterized.TestCase): def setUp(self): super(LayerCorrectnessTest, self).setUp() # Set two virtual CPUs to test MirroredStrategy with multiple devices cpus = config_module.list_physical_devices('CPU') config_module.set_logical_device_configuration(cpus[0], [ context.LogicalDeviceConfiguration(), context.LogicalDeviceConfiguration(), ]) def _create_model_from_layer(self, layer, input_shapes): inputs = [layers.Input(batch_input_shape=s) for s in input_shapes] if len(inputs) == 1: inputs = inputs[0] y = layer(inputs) model = models.Model(inputs, y) model.compile('sgd', 'mse') return model @parameterized.named_parameters( ('LeakyReLU', advanced_activations.LeakyReLU, (2, 2)), ('PReLU', advanced_activations.PReLU, (2, 2)), ('ELU', advanced_activations.ELU, (2, 2)), ('ThresholdedReLU', advanced_activations.ThresholdedReLU, (2, 2)), ('Softmax', advanced_activations.Softmax, (2, 2)), ('ReLU', advanced_activations.ReLU, (2, 2)), ('Conv1D', lambda: convolutional.Conv1D(2, 2), (2, 2, 1)), ('Conv2D', lambda: convolutional.Conv2D(2, 2), (2, 2, 2, 1)), ('Conv3D', lambda: convolutional.Conv3D(2, 2), (2, 2, 2, 2, 1)), ('Conv2DTranspose', lambda: convolutional.Conv2DTranspose(2, 2), (2, 2, 2, 2)), ('SeparableConv2D', lambda: convolutional.SeparableConv2D(2, 2), (2, 2, 2, 1)), ('DepthwiseConv2D', lambda: convolutional.DepthwiseConv2D(2, 2), (2, 2, 2, 1)), ('UpSampling2D', convolutional.UpSampling2D, (2, 2, 2, 1)), ('ZeroPadding2D', convolutional.ZeroPadding2D, (2, 2, 2, 1)), ('Cropping2D', convolutional.Cropping2D, (2, 3, 3, 1)), ('ConvLSTM2D', lambda: convolutional_recurrent.ConvLSTM2D(4, kernel_size=(2, 2)), (4, 4, 4, 4, 4)), ('Dense', lambda: core.Dense(2), (2, 2)), ('Dropout', lambda: core.Dropout(0.5), (2, 2)), ('SpatialDropout2D', lambda: core.SpatialDropout2D(0.5), (2, 2, 2, 2)), ('Activation', lambda: core.Activation('sigmoid'), (2, 2)), ('Reshape', lambda: core.Reshape((1, 4, 1)), (2, 2, 2)), ('Permute', lambda: core.Permute((2, 1)), (2, 2, 2)), ('Attention', dense_attention.Attention, [(2, 2, 3), (2, 3, 3), (2, 3, 3)]), ('AdditiveAttention', dense_attention.AdditiveAttention, [(2, 2, 3), (2, 3, 3), (2, 3, 3)]), ('Embedding', lambda: embeddings.Embedding(4, 4), (2, 4), 2e-3, 2e-3, np.random.randint(4, size=(2, 4))), ('LocallyConnected1D', lambda: local.LocallyConnected1D(2, 2), (2, 2, 1)), ('LocallyConnected2D', lambda: local.LocallyConnected2D(2, 2), (2, 2, 2, 1)), ('Add', merge.Add, [(2, 2), (2, 2)]), ('Subtract', merge.Subtract, [(2, 2), (2, 2)]), ('Multiply', merge.Multiply, [(2, 2), (2, 2)]), ('Average', merge.Average, [(2, 2), (2, 2)]), ('Maximum', merge.Maximum, [(2, 2), (2, 2)]), ('Minimum', merge.Minimum, [(2, 2), (2, 2)]), ('Concatenate', merge.Concatenate, [(2, 2), (2, 2)]), ('Dot', lambda: merge.Dot(1), [(2, 2), (2, 2)]), ('GaussianNoise', lambda: noise.GaussianNoise(0.5), (2, 2)), ('GaussianDropout', lambda: noise.GaussianDropout(0.5), (2, 2)), ('AlphaDropout', lambda: noise.AlphaDropout(0.5), (2, 2)), ('BatchNormalization', normalization_v2.BatchNormalization, (2, 2), 1e-2, 1e-2), ('LayerNormalization', normalization.LayerNormalization, (2, 2)), ('LayerNormalizationUnfused', lambda: normalization.LayerNormalization(axis=1), (2, 2, 2)), ('MaxPooling2D', pooling.MaxPooling2D, (2, 2, 2, 1)), ('AveragePooling2D', pooling.AveragePooling2D, (2, 2, 2, 1)), ('GlobalMaxPooling2D', pooling.GlobalMaxPooling2D, (2, 2, 2, 1)), ('GlobalAveragePooling2D', pooling.GlobalAveragePooling2D, (2, 2, 2, 1)), ('SimpleRNN', lambda: recurrent.SimpleRNN(units=4), (4, 4, 4), 1e-2, 1e-2), ('GRU', lambda: recurrent.GRU(units=4), (4, 4, 4)), ('LSTM', lambda: recurrent.LSTM(units=4), (4, 4, 4)), ('GRUV2', lambda: recurrent_v2.GRU(units=4), (4, 4, 4)), ('LSTMV2', lambda: recurrent_v2.LSTM(units=4), (4, 4, 4)), ('TimeDistributed', lambda: wrappers.TimeDistributed(core.Dense(2)), (2, 2, 2)), ('Bidirectional', lambda: wrappers.Bidirectional(recurrent.SimpleRNN(units=4)), (2, 2, 2)), ('AttentionLayerCausal', lambda: dense_attention.Attention(causal=True), [ (2, 2, 3), (2, 3, 3), (2, 3, 3) ]), ('AdditiveAttentionLayerCausal', lambda: dense_attention.AdditiveAttention(causal=True), [(2, 3, 4), (2, 3, 4), (2, 3, 4)]), ) def test_layer(self, f32_layer_fn, input_shape, rtol=2e-3, atol=2e-3, input_data=None): """Tests a layer by comparing the float32 and mixed precision weights. A float32 layer, a mixed precision layer, and a distributed mixed precision layer are run. The three layers are identical other than their dtypes and distribution strategies. The outputs after predict() and weights after fit() are asserted to be close. Args: f32_layer_fn: A function returning a float32 layer. The other two layers will automatically be created from this input_shape: The shape of the input to the layer, including the batch dimension. Or a list of shapes if the layer takes multiple inputs. rtol: The relative tolerance to be asserted. atol: The absolute tolerance to be asserted. input_data: A Numpy array with the data of the input. If None, input data will be randomly generated """ if f32_layer_fn == convolutional.ZeroPadding2D and \ test.is_built_with_rocm(): return if isinstance(input_shape[0], int): input_shapes = [input_shape] else: input_shapes = input_shape strategy = create_mirrored_strategy() f32_layer = f32_layer_fn() # Create the layers assert f32_layer.dtype == f32_layer._compute_dtype == 'float32' config = f32_layer.get_config() config['dtype'] = policy.Policy('mixed_float16') mp_layer = f32_layer.__class__.from_config(config) distributed_mp_layer = f32_layer.__class__.from_config(config) # Compute per_replica_input_shapes for the distributed model global_batch_size = input_shapes[0][0] assert global_batch_size % strategy.num_replicas_in_sync == 0, ( 'The number of replicas, %d, does not divide the global batch size of ' '%d' % (strategy.num_replicas_in_sync, global_batch_size)) per_replica_batch_size = ( global_batch_size // strategy.num_replicas_in_sync) per_replica_input_shapes = [(per_replica_batch_size,) + s[1:] for s in input_shapes] # Create the models f32_model = self._create_model_from_layer(f32_layer, input_shapes) mp_model = self._create_model_from_layer(mp_layer, input_shapes) with strategy.scope(): distributed_mp_model = self._create_model_from_layer( distributed_mp_layer, per_replica_input_shapes) # Set all model weights to the same values f32_weights = f32_model.get_weights() mp_model.set_weights(f32_weights) distributed_mp_model.set_weights(f32_weights) # Generate input data if input_data is None: # Cast inputs to float16 to avoid measuring error from having f16 layers # cast to float16. input_data = [np.random.normal(size=s).astype('float16') for s in input_shapes] if len(input_data) == 1: input_data = input_data[0] # Assert all models have close outputs. f32_output = f32_model.predict(input_data) mp_output = mp_model.predict(input_data) self.assertAllClose( mp_output, f32_output, rtol=rtol, atol=atol) self.assertAllClose( distributed_mp_model.predict(input_data), f32_output, rtol=rtol, atol=atol) # Run fit() on models output = np.random.normal(size=f32_model.outputs[0].shape).astype('float16') for model in f32_model, mp_model, distributed_mp_model: model.fit(input_data, output, batch_size=global_batch_size) # Assert all models have close weights f32_weights = f32_model.get_weights() self.assertAllClose( mp_model.get_weights(), f32_weights, rtol=rtol, atol=atol) self.assertAllClose( distributed_mp_model.get_weights(), f32_weights, rtol=rtol, atol=atol)