| # Copyright 2018 The TensorFlow Authors. All Rights Reserved. |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # ============================================================================== |
| """Test DistributionStrategy in the zero batch case.""" |
| |
| from absl.testing import parameterized |
| import numpy as np |
| |
| from tensorflow.python.data.ops import dataset_ops |
| from tensorflow.python.distribute import combinations |
| from tensorflow.python.distribute import strategy_combinations |
| from tensorflow.python.distribute import test_util |
| from tensorflow.python.eager import backprop |
| from tensorflow.python.eager import def_function |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.framework import ops |
| from tensorflow.python.layers import normalization |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import variables |
| from tensorflow.python.ops.losses import losses |
| from tensorflow.python.platform import test |
| from tensorflow.python.training import gradient_descent |
| |
| |
| class NormalizationTest(test.TestCase, parameterized.TestCase): |
| |
| @combinations.generate( |
| combinations.combine( |
| distribution=[ |
| strategy_combinations.one_device_strategy, |
| ], |
| mode=["graph"], |
| fused=[True, False])) |
| def testBNWithZeroBatchInputGraph(self, distribution, fused): |
| distribution.extended.experimental_enable_get_next_as_optional = True |
| with distribution.scope(), self.cached_session() as sess: |
| bn_list = [] |
| inputs = np.random.random((0, 4, 4, 3)) + 100 |
| targets = np.random.random((0, 4, 4, 3)) |
| inputs_placeholder = array_ops.placeholder( |
| dtype=dtypes.float32, shape=[None, 4, 4, 3]) |
| targets_placeholder = array_ops.placeholder( |
| dtype=dtypes.float32, shape=[None, 4, 4, 3]) |
| |
| def step_fn(is_training, inputs, targets=None): |
| bn = normalization.BatchNormalization( |
| axis=3, epsilon=1e-3, momentum=0.9, fused=fused) |
| bn_list.append(bn) |
| outputs = bn.apply(inputs, training=is_training) |
| if not is_training: |
| return outputs |
| |
| loss = losses.mean_squared_error(targets, outputs) |
| optimizer = gradient_descent.GradientDescentOptimizer(0.01) |
| train_op = optimizer.minimize(loss) |
| with ops.control_dependencies([train_op]): |
| return array_ops.identity(loss) |
| |
| train_op = distribution.extended.call_for_each_replica( |
| step_fn, args=(True, inputs_placeholder, targets_placeholder)) |
| predict_op = distribution.extended.call_for_each_replica( |
| step_fn, args=(False, inputs_placeholder)) |
| bn = bn_list[0] |
| |
| self.evaluate(variables.global_variables_initializer()) |
| |
| # Check for initial statistics and weights. |
| moving_mean, moving_var = self.evaluate( |
| [bn.moving_mean, bn.moving_variance]) |
| self.assertAllEqual([0, 0, 0], moving_mean) |
| self.assertAllEqual([1, 1, 1], moving_var) |
| |
| np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta]) |
| self.assertAllEqual([1, 1, 1], np_gamma) |
| self.assertAllEqual([0, 0, 0], np_beta) |
| |
| for _ in range(100): |
| np_output, _, _ = sess.run([train_op] + bn.updates, { |
| inputs_placeholder: inputs, |
| targets_placeholder: targets |
| }) |
| self.assertEqual(0.0, np_output) |
| |
| # Verify that the statistics and weights are not changed after training. |
| moving_mean, moving_var = self.evaluate( |
| [bn.moving_mean, bn.moving_variance]) |
| self.assertAllEqual([0, 0, 0], moving_mean) |
| self.assertAllEqual([1, 1, 1], moving_var) |
| |
| np_gamma, np_beta = self.evaluate([bn.gamma, bn.beta]) |
| self.assertAllEqual([1, 1, 1], np_gamma) |
| self.assertAllEqual([0, 0, 0], np_beta) |
| |
| # Test inference. |
| np_output = sess.run(predict_op, {inputs_placeholder: inputs}) |
| self.assertEqual([], np_output.tolist()) |
| |
| @combinations.generate( |
| combinations.combine( |
| distribution=[ |
| strategy_combinations.one_device_strategy, |
| ], |
| mode=["eager"], |
| fused=[True, False])) |
| def testBNWithZeroBatchInput(self, distribution, fused): |
| distribution.extended.experimental_enable_get_next_as_optional = True |
| with distribution.scope(): |
| inputs = np.random.random((0, 4, 4, 3)).astype(np.float32) + 100 |
| targets = np.random.random((0, 4, 4, 3)).astype(np.float32) |
| bn = normalization.BatchNormalization( |
| axis=3, epsilon=1e-3, momentum=0.9, fused=fused) |
| optimizer = gradient_descent.GradientDescentOptimizer(0.01) |
| |
| @def_function.function |
| def train_step(): |
| def step_fn(inputs, targets): |
| with backprop.GradientTape() as tape: |
| outputs = bn.apply(inputs, training=True) |
| loss = losses.mean_squared_error(targets, outputs) |
| grads = tape.gradient(loss, bn.variables) |
| optimizer.apply_gradients(zip(grads, bn.variables)) |
| return loss |
| |
| return distribution.run(step_fn, args=(inputs, targets)) |
| |
| for _ in range(100): |
| np_output = train_step().numpy() |
| self.assertEqual(0.0, np_output) |
| |
| # Verify that the statistics and weights are not changed after training. |
| self.assertAllEqual([0, 0, 0], bn.moving_mean.numpy()) |
| self.assertAllEqual([1, 1, 1], bn.moving_variance.numpy()) |
| self.assertAllEqual([1, 1, 1], bn.gamma.numpy()) |
| self.assertAllEqual([0, 0, 0], bn.beta.numpy()) |
| |
| @def_function.function |
| def test_step(): |
| def step_fn(inputs): |
| outputs = bn.apply(inputs, training=False) |
| return outputs |
| |
| return distribution.run(step_fn, args=(inputs,)) |
| |
| # Test inference. |
| self.assertAllEqual(np.zeros(shape=(0, 4, 4, 3), dtype=np.float32), |
| test_step().numpy()) |
| |
| @combinations.generate( |
| combinations.combine( |
| distribution=[ |
| strategy_combinations.one_device_strategy, |
| ], |
| mode=["eager"], |
| fused=[True, False])) |
| def testBNWithDynamicBatchInputEager(self, distribution, fused): |
| distribution.extended.experimental_enable_get_next_as_optional = True |
| with distribution.scope(): |
| # Explicitly create dataset with drop_remainder=False. |
| # This would make batch size unknown. |
| inputs = np.random.random((11, 4, 4, 3)).astype(np.float32) + 100 |
| targets = np.random.random((11, 4, 4, 3)).astype(np.float32) |
| dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)).batch( |
| 10, drop_remainder=False).repeat() |
| dataset_iterator = iter( |
| distribution.experimental_distribute_dataset(dataset)) |
| |
| bn = normalization.BatchNormalization( |
| axis=-1, epsilon=1e-3, momentum=0.9, fused=fused) |
| optimizer = gradient_descent.GradientDescentOptimizer(0.01) |
| |
| @def_function.function |
| def train_step(iterator): |
| |
| def step_fn(inputs): |
| features, targets = inputs |
| with backprop.GradientTape() as tape: |
| outputs = bn(features, training=True) |
| loss = losses.mean_squared_error(targets, outputs) |
| |
| grads = tape.gradient(loss, bn.variables) |
| optimizer.apply_gradients(zip(grads, bn.variables)) |
| return loss |
| |
| return distribution.run(step_fn, args=(next(iterator),)) |
| |
| for _ in range(100): |
| train_step(dataset_iterator).numpy() |
| |
| # Verify that the statistics and weights are updated. |
| self.assertNotAllEqual(np.ndarray([0, 0, 0]), bn.moving_mean.numpy()) |
| self.assertNotAllEqual(np.ndarray([1, 1, 1]), bn.moving_variance.numpy()) |
| self.assertNotAllEqual(np.ndarray([1, 1, 1]), bn.gamma.numpy()) |
| self.assertNotAllEqual(np.ndarray([0, 0, 0]), bn.beta.numpy()) |
| |
| |
| if __name__ == "__main__": |
| test_util.main() |
