zero_batch_test.py

# 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()