wehub-resource-sync
209 lines
8.1 KiB
Python
209 lines
8.1 KiB
Python
# 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()
|