wehub-resource-sync
1581 lines
53 KiB
Python
1581 lines
53 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.
|
|
# ==============================================================================
|
|
"""Tests for TPUStrategy."""
|
|
|
|
from absl.testing import parameterized
|
|
|
|
from tensorflow.core.protobuf import config_pb2
|
|
from tensorflow.python import tf2
|
|
from tensorflow.python.data.ops import dataset_ops
|
|
from tensorflow.python.distribute import distribute_lib
|
|
from tensorflow.python.distribute import reduce_util
|
|
from tensorflow.python.distribute import strategy_test_lib
|
|
from tensorflow.python.distribute import tpu_strategy as tpu_lib
|
|
from tensorflow.python.distribute import tpu_values
|
|
from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
|
|
from tensorflow.python.eager import def_function
|
|
from tensorflow.python.eager import remote
|
|
from tensorflow.python.eager import test
|
|
from tensorflow.python.framework import composite_tensor
|
|
from tensorflow.python.framework import config
|
|
from tensorflow.python.framework import constant_op
|
|
from tensorflow.python.framework import device as tf_device
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import errors
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import sparse_tensor
|
|
from tensorflow.python.framework import tensor_spec
|
|
from tensorflow.python.framework import test_util
|
|
from tensorflow.python.framework import type_spec
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import control_flow_switch_case
|
|
from tensorflow.python.ops import embedding_ops
|
|
from tensorflow.python.ops import gen_dataset_ops
|
|
from tensorflow.python.ops import lookup_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops import random_ops
|
|
from tensorflow.python.ops import variable_scope
|
|
from tensorflow.python.ops import variables
|
|
from tensorflow.python.ops.ragged import ragged_tensor
|
|
from tensorflow.python.platform import flags
|
|
from tensorflow.python.platform import tf_logging as logging
|
|
from tensorflow.python.saved_model import load
|
|
from tensorflow.python.saved_model import save
|
|
from tensorflow.python.tpu import device_assignment as device_assignment_lib
|
|
from tensorflow.python.tpu import tpu_hardware_feature
|
|
from tensorflow.python.tpu import tpu_replication
|
|
from tensorflow.python.tpu import tpu_strategy_util
|
|
from tensorflow.python.trackable import autotrackable
|
|
from tensorflow.python.training import server_lib
|
|
from tensorflow.python.util import nest
|
|
|
|
FLAGS = flags.FLAGS
|
|
flags.DEFINE_string("tpu", "", "Name of TPU to connect to.")
|
|
flags.DEFINE_string("project", None, "Name of GCP project with TPU.")
|
|
flags.DEFINE_string("zone", None, "Name of GCP zone with TPU.")
|
|
|
|
|
|
class TestExportArchive(autotrackable.AutoTrackable):
|
|
|
|
def __init__(self, init_value):
|
|
self._var = variables.Variable(init_value)
|
|
self._fn = def_function.function(self.update_var)
|
|
self._packed_var = self._var._packed_var
|
|
|
|
@def_function.function
|
|
def update_var(self):
|
|
# Use packed variable in function to simulate the error
|
|
# in b/323080532
|
|
self._packed_var.assign_add(3.0).assign_sub(1.0)
|
|
|
|
def save_function(self, directory):
|
|
save.save(self, directory)
|
|
|
|
@property
|
|
def packed_var(self):
|
|
return self._packed_var
|
|
|
|
|
|
def get_tpu_cluster_resolver():
|
|
resolver = tpu_cluster_resolver.TPUClusterResolver(
|
|
tpu=FLAGS.tpu,
|
|
zone=FLAGS.zone,
|
|
project=FLAGS.project,
|
|
)
|
|
return resolver
|
|
|
|
|
|
def get_tpu_strategy(enable_packed_var=False):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
strategy = tpu_lib.TPUStrategyV2(resolver)
|
|
strategy._enable_packed_variable_in_eager_mode = enable_packed_var
|
|
return strategy
|
|
|
|
|
|
# TPU tests which don't use TPUStrategy.
|
|
@test_util.with_eager_op_as_function
|
|
class TPUTest(test.TestCase):
|
|
|
|
# In this case, the entire computation in foo is compiled using JIT
|
|
# compilation.
|
|
def test_single_tpu_jit_compile(self):
|
|
with ops.device("/device:TPU:0"):
|
|
a = variables.Variable(1)
|
|
|
|
def get_a_plus_one():
|
|
return a + 1
|
|
|
|
@def_function.function(
|
|
input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
|
|
def foo(x):
|
|
b = x + get_a_plus_one()
|
|
b = b + get_a_plus_one()
|
|
return b + 1
|
|
|
|
with ops.device("/device:TPU:0"):
|
|
result = foo(a)
|
|
self.assertAllEqual(6, result)
|
|
|
|
# In this case, each of the ops in the TPU device scope are compiled and run
|
|
# individually.
|
|
def test_single_tpu_on_demand(self):
|
|
with ops.device("/device:TPU:0"):
|
|
a = variables.Variable(1)
|
|
|
|
def get_a_plus_one():
|
|
return a + 1
|
|
|
|
x = 1
|
|
with ops.device("/device:TPU:0"):
|
|
b = x + get_a_plus_one()
|
|
b = b + get_a_plus_one()
|
|
result = b + 1
|
|
|
|
self.assertAllEqual(6, result)
|
|
|
|
# In this case, each of the ops in the tf.function and TPU device scope are
|
|
# compiled and run individually.
|
|
def test_single_tpu_on_demand_tf_function(self):
|
|
with ops.device("/device:TPU:0"):
|
|
a = variables.Variable(1)
|
|
|
|
def get_a_plus_one():
|
|
return a + 1
|
|
|
|
@def_function.function(
|
|
input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
|
|
def foo(x):
|
|
with ops.device("/device:TPU:0"):
|
|
b = x + get_a_plus_one()
|
|
b = b + get_a_plus_one()
|
|
return b + 1
|
|
|
|
result = foo(a)
|
|
self.assertAllEqual(6, result)
|
|
|
|
def test_multiple_initialize_system(self):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
|
|
with test.mock.patch.object(logging, "warning") as mock_log:
|
|
tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
self.assertRegex(str(mock_log.call_args), "already been initialized")
|
|
|
|
def test_initialize_tpu_system_impl_input(self):
|
|
resolver = get_tpu_cluster_resolver()
|
|
with self.assertRaisesRegex(
|
|
TypeError,
|
|
r"tpu_cluster_resolver_cls is not"
|
|
r" tf.distribute.cluster_resolver.TPUClusterResolver."):
|
|
tpu_strategy_util.initialize_tpu_system_impl(
|
|
resolver, tpu_cluster_resolver_cls=None)
|
|
|
|
def test_shutdown_tpu_system_impl_input(self):
|
|
resolver = get_tpu_cluster_resolver()
|
|
with self.assertRaisesRegex(
|
|
TypeError,
|
|
r"tpu_cluster_resolver_cls is not"
|
|
r" tf.distribute.cluster_resolver.TPUClusterResolver."):
|
|
tpu_strategy_util.shutdown_tpu_system_impl(
|
|
resolver, tpu_cluster_resolver_cls=None)
|
|
|
|
def test_tpu_tf_function_same_device(self):
|
|
with ops.device("/device:TPU:0"):
|
|
a = variables.Variable(1)
|
|
|
|
@def_function.function(experimental_attributes={"_noinline": True})
|
|
def get_a_plus_one():
|
|
return a + 1
|
|
|
|
@def_function.function(
|
|
input_signature=[tensor_spec.TensorSpec([], dtypes.int32)])
|
|
def foo(x):
|
|
with ops.device("/device:TPU:0"):
|
|
b = x + get_a_plus_one()
|
|
return b + 1
|
|
|
|
result = foo(a)
|
|
self.assertAllEqual(4, result)
|
|
|
|
def test_tpu_return_int32(self):
|
|
with ops.device("/device:TPU:0"):
|
|
a = variables.Variable(0)
|
|
|
|
@def_function.function
|
|
def foo():
|
|
return a + 1
|
|
|
|
@def_function.function
|
|
def bar():
|
|
with ops.device("/device:TPU:1"):
|
|
return foo()
|
|
|
|
with ops.device("/device:CPU:0"):
|
|
result = bar() + 1
|
|
self.assertAllEqual(result, 2)
|
|
|
|
def test_tpu_output_device(self):
|
|
|
|
def foo():
|
|
return 1 + 1
|
|
|
|
func1 = def_function.function(foo, jit_compile=False)
|
|
func2 = def_function.function(
|
|
foo,
|
|
jit_compile=False,
|
|
experimental_attributes={
|
|
"_OutputsOnOpDevice": True,
|
|
},
|
|
)
|
|
|
|
with ops.device("/device:TPU:0"):
|
|
ret1 = func1()
|
|
ret2 = func2()
|
|
|
|
self.assertAllEqual(ret1.backing_device,
|
|
"/job:localhost/replica:0/task:0/device:CPU:0")
|
|
self.assertAllEqual(ret2.backing_device,
|
|
"/job:localhost/replica:0/task:0/device:TPU:0")
|
|
|
|
def test_on_demand_op_with_dynamic_output(self):
|
|
with ops.device("/device:TPU:0"):
|
|
where_output = array_ops.where([True, False, True])
|
|
self.assertAllEqual(where_output, [[0], [2]])
|
|
|
|
with ops.device("/device:TPU:0"):
|
|
repeat_output = array_ops.repeat(math_ops.range(2), [1, 4])
|
|
self.assertAllEqual(repeat_output, [0, 1, 1, 1, 1])
|
|
|
|
|
|
@parameterized.named_parameters([("PackedVar", True), ("", False)])
|
|
@test_util.with_eager_op_as_function
|
|
class TPUStrategyTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def test_handle_in_cross_replica_context(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
with strategy.scope():
|
|
v = variables.Variable(1.0)
|
|
|
|
@def_function.function
|
|
def func():
|
|
self.assertEndsWith(v.handle.device, "device:TPU:0")
|
|
return v + 1.0
|
|
|
|
ret = func()
|
|
self.assertAllEqual(ret, 2.0)
|
|
|
|
def test_save(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
with strategy.scope():
|
|
v = variables.Variable(1.0)
|
|
export_dir = self.create_tempdir()
|
|
save.save(v, export_dir)
|
|
reloaded_var = load.load(export_dir)
|
|
self.assertAllEqual(reloaded_var, 1.0)
|
|
|
|
def test_packed_variable_export(self, enable_packed_var):
|
|
if not enable_packed_var:
|
|
self.skipTest("Test for Packed Variables only.")
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
with strategy.scope():
|
|
export_dir = self.get_temp_dir()
|
|
export_archive = TestExportArchive(1.0)
|
|
export_archive.save_function(export_dir)
|
|
restored_object = load.load(export_dir)
|
|
with ops.device("/tpu:0"):
|
|
self.assertAllEqual(restored_object._packed_var, 1.0)
|
|
|
|
def testStaticHashTableDatasetFnHostTrainingLoop(self, enable_packed_var):
|
|
self._dataset_fn_tracing_count = 0
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
with strategy.scope():
|
|
vals = [0, 1, 2]
|
|
keys_tensor = constant_op.constant(
|
|
list(range(len(vals))), dtype=dtypes.int64)
|
|
vals_tensor = constant_op.constant(vals)
|
|
initializer = lookup_ops.KeyValueTensorInitializer(
|
|
keys_tensor, vals_tensor)
|
|
per_worker_table = lookup_ops.StaticHashTable(
|
|
initializer, default_value=-1)
|
|
|
|
@def_function.function
|
|
def dataset_fn(input_context):
|
|
tensor = constant_op.constant([0, 1, 3], dtype=dtypes.int64)
|
|
global_batch_size = 2
|
|
batch_size = input_context.get_per_replica_batch_size(global_batch_size)
|
|
dataset = dataset_ops.Dataset.from_tensors(tensor).repeat().batch(
|
|
batch_size, drop_remainder=True)
|
|
dataset = dataset.shard(input_context.num_input_pipelines,
|
|
input_context.input_pipeline_id)
|
|
dataset = dataset.prefetch(2) # This prefetches 2 batches per device.
|
|
dataset = dataset.map(per_worker_table.lookup)
|
|
self._dataset_fn_tracing_count += 1
|
|
return dataset
|
|
|
|
dist_iterator = iter(
|
|
strategy.experimental_distribute_datasets_from_function(dataset_fn))
|
|
|
|
@def_function.function
|
|
def step_fn(inputs):
|
|
# inputs should be [0, 1, -1]
|
|
return math_ops.reduce_sum(inputs)
|
|
|
|
def train_steps(iterator, steps):
|
|
|
|
for _ in math_ops.range(steps):
|
|
strategy.run(step_fn, args=(next(iterator),))
|
|
|
|
train_steps(dist_iterator, steps=5)
|
|
self.assertEqual(self._dataset_fn_tracing_count, 1)
|
|
|
|
def test_function_compile_with_xla(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
with strategy.scope():
|
|
v = variables.Variable(1.0)
|
|
|
|
@def_function.function
|
|
def func():
|
|
return v.read_value() + 1.0
|
|
|
|
with ops.device("/device:TPU:0"):
|
|
self.assertAllEqual(func(), 2.0)
|
|
|
|
def test_sequential_runs(self, enable_packed_var):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
# Computation replicated to all cores.
|
|
device_assignment = device_assignment_lib.DeviceAssignment.build(
|
|
topology, num_replicas=2)
|
|
strategy = tpu_lib.TPUStrategyV2(
|
|
resolver, experimental_device_assignment=device_assignment)
|
|
strategy._enable_packed_variable_in_eager_mode = enable_packed_var
|
|
|
|
# Computation on the 1st core.
|
|
device_assignment2 = device_assignment_lib.DeviceAssignment.build(
|
|
topology, num_replicas=1)
|
|
strategy2 = tpu_lib.TPUStrategyV2(
|
|
resolver, experimental_device_assignment=device_assignment2)
|
|
|
|
def computation(x):
|
|
return math_ops.square(x)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
outputs = strategy.experimental_local_results(
|
|
strategy.run(computation, args=([2., 2.],)))
|
|
outputs2 = strategy2.run(
|
|
computation, args=([outputs[0]],))
|
|
return outputs2
|
|
|
|
self.assertAllEqual([[16., 16.]], train_step())
|
|
|
|
def test_device_switch_case(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
with strategy.scope():
|
|
a = variables.Variable(1)
|
|
|
|
inference_iteration = variables.Variable(-1)
|
|
|
|
def inference_fn(x, i):
|
|
return a + x + i
|
|
|
|
@def_function.function
|
|
def run_inference(x):
|
|
|
|
def do_inference(device, inference_fn, i):
|
|
with ops.device(device):
|
|
return inference_fn(x, i)
|
|
|
|
branch_fns = {
|
|
0: (lambda: do_inference("/device:TPU:0", inference_fn, 0)),
|
|
1: (lambda: do_inference("/device:TPU:1", inference_fn, 1)),
|
|
}
|
|
branch_index = inference_iteration.assign_add(1, use_locking=True) % 2
|
|
return control_flow_switch_case.switch_case(branch_index, branch_fns)
|
|
|
|
self.assertAllEqual(2., run_inference(1)) # Use TPU core 0.
|
|
self.assertAllEqual(3., run_inference(1)) # Use TPU core 1.
|
|
|
|
def test_recover_from_compilation_failures(self, enable_packed_var):
|
|
# TODO(b/148150981): Stop skipping this test once recovery works
|
|
# for non-local TPU.
|
|
if FLAGS.tpu:
|
|
self.skipTest("Recovery fails for non-local TPU, see b/148150981")
|
|
|
|
# Disable automatic outside compilation.
|
|
config.set_soft_device_placement(False)
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
@def_function.function
|
|
def compilation_failure_run():
|
|
|
|
def computation():
|
|
return random_ops.random_gamma([10], [0.5, 1.5])
|
|
|
|
return strategy.run(computation)
|
|
|
|
with self.assertRaises(errors.OpError):
|
|
compilation_failure_run()
|
|
|
|
@def_function.function
|
|
def good_run():
|
|
|
|
def computation():
|
|
return random_ops.random_normal([10])
|
|
|
|
return strategy.run(computation)
|
|
|
|
good_run()
|
|
|
|
def test_dynamic_shape_with_outside_compilation_failure(
|
|
self, enable_packed_var):
|
|
# Enable automatic outside compilation.
|
|
config.set_soft_device_placement(True)
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
dataset = dataset_ops.Dataset.from_tensors(("string", 1.0)).repeat().batch(
|
|
2, drop_remainder=False)
|
|
dataset = strategy.experimental_distribute_dataset(dataset)
|
|
iterator = iter(dataset)
|
|
|
|
@def_function.function
|
|
def train_fn(iterator):
|
|
|
|
def step_fn(inputs):
|
|
input0, input1 = inputs
|
|
return array_ops.size(input0), math_ops.reduce_sum(input1)
|
|
|
|
return strategy.experimental_local_results(
|
|
strategy.run(step_fn, args=(next(iterator),)))
|
|
|
|
with self.assertRaises(errors.InvalidArgumentError):
|
|
logging.info(train_fn(iterator))
|
|
|
|
def test_computation_on_subset_cores(self, enable_packed_var):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
all_core_strategy = tpu_lib.TPUStrategyV2(resolver)
|
|
all_core_strategy._enable_packed_variable_in_eager_mode = enable_packed_var
|
|
|
|
with all_core_strategy.scope():
|
|
v = variables.Variable(0.0,
|
|
aggregation=variables.VariableAggregation.MEAN)
|
|
|
|
# Computation on the 1st core.
|
|
device_assignment = device_assignment_lib.DeviceAssignment.build(
|
|
topology, num_replicas=1)
|
|
first_core_strategy = tpu_lib.TPUStrategyV2(
|
|
resolver, experimental_device_assignment=device_assignment)
|
|
first_core_strategy._enable_packed_variable_in_eager_mode = (
|
|
enable_packed_var)
|
|
|
|
# Computation on the 2nd core.
|
|
device_assignment2 = device_assignment_lib.DeviceAssignment(
|
|
topology, [[[0, 0, 0, 1]]])
|
|
second_core_strategy = tpu_lib.TPUStrategyV2(
|
|
resolver, experimental_device_assignment=device_assignment2)
|
|
second_core_strategy._enable_packed_variable_in_eager_mode = (
|
|
enable_packed_var)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
|
|
def step_fn():
|
|
return v + 1.0
|
|
|
|
all_core_strategy.run(step_fn)
|
|
r1 = first_core_strategy.run(step_fn)
|
|
r2 = second_core_strategy.run(step_fn)
|
|
return r1 + r2
|
|
|
|
train_step()
|
|
self.assertAllEqual(2., train_step())
|
|
|
|
def test_worker_devices_on_subset_cores(self, enable_packed_var):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
|
|
# Strategy for the 1st core.
|
|
device_assignment = device_assignment_lib.DeviceAssignment.build(
|
|
topology, num_replicas=1)
|
|
first_core_strategy = tpu_lib.TPUStrategyV2(
|
|
resolver, experimental_device_assignment=device_assignment)
|
|
first_core_strategy._enable_packed_variable_in_eager_mode = (
|
|
enable_packed_var)
|
|
|
|
# Strategy for the 2nd core.
|
|
device_assignment2 = device_assignment_lib.DeviceAssignment(
|
|
topology, [[[0, 0, 0, 1]]])
|
|
second_core_strategy = tpu_lib.TPUStrategyV2(
|
|
resolver, experimental_device_assignment=device_assignment2)
|
|
second_core_strategy._enable_packed_variable_in_eager_mode = (
|
|
enable_packed_var)
|
|
|
|
self.assertLen(first_core_strategy.extended.worker_devices, 1)
|
|
self.assertEndsWith(first_core_strategy.extended.worker_devices[0],
|
|
"device:TPU:0")
|
|
|
|
self.assertLen(second_core_strategy.extended.worker_devices, 1)
|
|
self.assertEndsWith(second_core_strategy.extended.worker_devices[0],
|
|
"device:TPU:1")
|
|
|
|
def test_control_output_in_while_body_fn(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
with strategy.scope():
|
|
v = variables.Variable(
|
|
0.0, aggregation=variables.VariableAggregation.MEAN)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
|
|
def step_fn():
|
|
v.assign_add(1)
|
|
|
|
for _ in math_ops.range(2):
|
|
strategy.run(step_fn)
|
|
|
|
train_step()
|
|
self.assertEqual(2.0, v.numpy())
|
|
|
|
def test_cluster_conditional_with_dynamic_shape(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
|
|
def shape_list(tensor):
|
|
shape = tensor.shape.as_list()
|
|
|
|
non_static_indexes = []
|
|
for (index, dim) in enumerate(shape):
|
|
if dim is None:
|
|
non_static_indexes.append(index)
|
|
|
|
if not non_static_indexes:
|
|
return shape
|
|
|
|
dynamic_shape = array_ops.shape(input=tensor)
|
|
for index in non_static_indexes:
|
|
shape[index] = dynamic_shape[index]
|
|
|
|
return shape
|
|
|
|
def step_fn(condition):
|
|
where = array_ops.where(condition)
|
|
if array_ops.shape(where)[0] > 0:
|
|
tensor_shape = shape_list(where)
|
|
d1 = tensor_shape[0]
|
|
d2 = tensor_shape[1]
|
|
where = array_ops.reshape(where, [d1, d2])
|
|
return where
|
|
|
|
return strategy.run(step_fn, args=([True, False, True],))
|
|
|
|
outputs = strategy.experimental_local_results(train_step())
|
|
self.assertAllEqual(outputs[0].numpy(), [[0], [2]])
|
|
|
|
def test_cluster_in_graph_and_while_body_fn(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
|
|
def step_fn(prev):
|
|
s = prev + 1
|
|
return s
|
|
|
|
def init_fn():
|
|
return array_ops.zeros(shape=())
|
|
|
|
prev = strategy.run(init_fn)
|
|
for _ in math_ops.range(10):
|
|
prev = strategy.run(step_fn, args=(prev,))
|
|
return strategy.reduce(reduce_util.ReduceOp.SUM, prev, axis=None)
|
|
|
|
sum_val = train_step().numpy().astype(float)
|
|
self.assertEqual(sum_val, strategy.num_replicas_in_sync * 10)
|
|
|
|
def test_two_clusters_with_same_fn(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
@def_function.function
|
|
def foo(x):
|
|
return strategy.run(lambda x: x + 1, (x,))
|
|
|
|
@def_function.function
|
|
def bar(x):
|
|
foo(x)
|
|
return foo(x)
|
|
|
|
bar(1)
|
|
|
|
def test_tpu_variable_run_argument(self, enable_packed_var):
|
|
# TPUStrategy.run() casts inputs to Tensor, but has logic to preserve
|
|
# variables to avoid unintuitive errors.
|
|
# Here we test that a TPUDistributedVariable passed to TPUStrategy.run()
|
|
# remains a variable.
|
|
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
with strategy.scope():
|
|
tpu_variable = variables.Variable(1)
|
|
|
|
def replica_step(first_arg, variable):
|
|
del first_arg # Just here to make sure we're not relying on arg position.
|
|
|
|
if variable is not None:
|
|
self.assertIsInstance(variable, tpu_values.TPUDistributedVariable)
|
|
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(
|
|
replica_step, args=(
|
|
2,
|
|
tpu_variable,
|
|
))
|
|
|
|
step()
|
|
|
|
def test_tpu_run_arg_parsing(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
with strategy.scope():
|
|
tpu_vars = [variables.Variable(1)]
|
|
|
|
def only_star_args(*args):
|
|
del args
|
|
|
|
def pos_and_star_args(first_arg, *args):
|
|
del first_arg
|
|
del args
|
|
|
|
def named_args(first_arg, second_arg):
|
|
del first_arg
|
|
del second_arg
|
|
|
|
def star_args_and_kw_only(*args, kw):
|
|
del args
|
|
del kw
|
|
|
|
# pylint:disable=function-redefined
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(only_star_args, args=(2,))
|
|
|
|
step()
|
|
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(named_args, kwargs={"first_arg": 2, "second_arg": 3})
|
|
|
|
step()
|
|
|
|
with self.assertRaisesRegex(TypeError, r"got multiple values for argument"):
|
|
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(
|
|
named_args, args=(1,), kwargs={
|
|
"first_arg": 2,
|
|
"second_arg": 3
|
|
})
|
|
|
|
step()
|
|
|
|
with self.assertRaisesRegex(ValueError,
|
|
r"cannot handle Variables passed to \*args"):
|
|
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(
|
|
only_star_args, args=(
|
|
2,
|
|
tpu_vars,
|
|
))
|
|
|
|
step()
|
|
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(pos_and_star_args, args=(2, 3, 4))
|
|
|
|
step()
|
|
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(star_args_and_kw_only, args=(2, 3), kwargs={"kw": tpu_vars})
|
|
|
|
step()
|
|
|
|
with self.assertRaisesRegex(ValueError,
|
|
r"mix of positional args and \*args"):
|
|
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(pos_and_star_args, args=(tpu_vars, 3, 4))
|
|
|
|
step()
|
|
|
|
with self.assertRaisesRegex(ValueError, r"Too many positional arguments"):
|
|
|
|
@def_function.function
|
|
def step():
|
|
strategy.run(named_args, args=(2, 3, 4))
|
|
|
|
step()
|
|
|
|
class DummyClass:
|
|
|
|
@def_function.function
|
|
def method(self, arg_1):
|
|
del arg_1
|
|
|
|
def step(self):
|
|
strategy.run(self.method, args=(tpu_vars,))
|
|
|
|
DummyClass().step()
|
|
# pylint:enable=function-redefined
|
|
|
|
def test_using_external_variable_inside_tf_function(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
dataset = dataset_ops.Dataset.range(
|
|
strategy.num_replicas_in_sync * 2,
|
|
output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
|
|
input_iterator = iter(strategy.experimental_distribute_dataset(dataset))
|
|
|
|
v = variables.Variable(2.0)
|
|
|
|
@def_function.function
|
|
def train_step(data):
|
|
def computation(inputs):
|
|
return inputs + v
|
|
return strategy.run(computation, args=(data,))
|
|
|
|
expected_result = [[x + 2.] for x in range(0, strategy.num_replicas_in_sync)
|
|
]
|
|
self.assertAllEqual(
|
|
expected_result,
|
|
strategy.experimental_local_results(train_step(next(input_iterator))))
|
|
|
|
# TODO(b/145574622): Remove this test once it is re-enabled in values_test.py.
|
|
def test_all_reduce_on_sync_on_read_variable(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
dataset = dataset_ops.Dataset.range(
|
|
strategy.num_replicas_in_sync, output_type=dtypes.float32).batch(
|
|
strategy.num_replicas_in_sync, drop_remainder=True)
|
|
input_iterator = iter(strategy.experimental_distribute_dataset(dataset))
|
|
|
|
with strategy.scope():
|
|
w = variables.Variable(
|
|
(0.,),
|
|
shape=(1,),
|
|
trainable=False,
|
|
synchronization=variables.VariableSynchronization.ON_READ,
|
|
aggregation=variables.VariableAggregation.ONLY_FIRST_REPLICA)
|
|
|
|
self.assertFalse(w._is_mirrored())
|
|
|
|
@def_function.function
|
|
def run(iterator):
|
|
|
|
def computation(x):
|
|
w.assign(x + w)
|
|
return w
|
|
|
|
def all_reduce(x):
|
|
ctx = distribute_lib.get_replica_context()
|
|
return ctx.all_reduce("SUM", w) + x
|
|
|
|
outputs = strategy.run(computation, args=(next(iterator),))
|
|
outputs2 = strategy.experimental_local_results(
|
|
strategy.run(all_reduce, args=(outputs,)))
|
|
return outputs2
|
|
|
|
data = range(0, strategy.num_replicas_in_sync)
|
|
data_sum = sum(data)
|
|
expected_result = [
|
|
[x + data_sum] for x in range(0, strategy.num_replicas_in_sync)
|
|
]
|
|
self.assertAllEqual(expected_result, run(input_iterator))
|
|
self.assertAllEqual((0.,), w.read_value())
|
|
|
|
def test_run_output_on_device(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
def computation(x):
|
|
return math_ops.square(x)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
outputs = strategy.experimental_local_results(
|
|
strategy.run(computation, args=(2,)))
|
|
return outputs
|
|
|
|
results = train_step()
|
|
self.assertAllEqual([4., 4.], results)
|
|
self.assertAllEqual("/job:localhost/replica:0/task:0/device:TPU:0",
|
|
results[0].backing_device)
|
|
self.assertAllEqual("/job:localhost/replica:0/task:0/device:TPU:1",
|
|
results[1].backing_device)
|
|
|
|
def test_run_passing_and_returning_nones(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
|
|
def computation(x):
|
|
return x
|
|
|
|
# Note that this input None is nested.
|
|
outputs = strategy.experimental_local_results(
|
|
strategy.run(computation, args=([1, [2, None]],)))
|
|
return outputs
|
|
|
|
results = train_step()
|
|
|
|
self.assertAllEqual(1, results[0][0])
|
|
self.assertAllEqual(2, results[0][1][0])
|
|
self.assertIsNone(results[0][1][1])
|
|
|
|
def test_run_passing_and_returning_empty_list(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
|
|
def computation(x):
|
|
return x
|
|
|
|
outputs = strategy.experimental_local_results(
|
|
strategy.run(computation, args=([],)))
|
|
return outputs
|
|
|
|
self.assertEqual([], train_step()[0])
|
|
|
|
def test_run_passing_and_returning_empty_dict(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
@def_function.function
|
|
def train_step():
|
|
|
|
def computation(x):
|
|
return x
|
|
|
|
outputs = strategy.experimental_local_results(
|
|
strategy.run(computation, args=({},)))
|
|
return outputs
|
|
|
|
self.assertEqual({}, train_step()[0])
|
|
|
|
def test_composite_input_output(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
if strategy.num_replicas_in_sync != 2:
|
|
self.skipTest("Test assumes two replicas.")
|
|
|
|
with strategy.scope():
|
|
table = variables.Variable(
|
|
initial_value=[[0.0, 1.0], [3.0, 7.0]], dtype=dtypes.float32)
|
|
|
|
@def_function.function
|
|
def sparse_lookup(iterator):
|
|
|
|
def tpu_function(sparse):
|
|
# Assumes dense_shape is (2, *)
|
|
looked_up = array_ops.gather(table, sparse.values)
|
|
segment_sum = math_ops.unsorted_segment_sum(
|
|
looked_up, sparse.indices[:, 0], 2)
|
|
return sparse, segment_sum
|
|
|
|
return nest.map_structure(
|
|
strategy.experimental_local_results,
|
|
strategy.run(tpu_function, args=(next(iterator),)))
|
|
|
|
def dataset_fn(_):
|
|
dataset = dataset_ops.Dataset.range(2)
|
|
|
|
def make_sparse(_):
|
|
return sparse_tensor.SparseTensor(
|
|
indices=array_ops.constant([[0, 0], [1, 0], [1, 1]],
|
|
dtype=dtypes.int64),
|
|
values=array_ops.constant([0, 0, 1], dtype=dtypes.int32),
|
|
dense_shape=array_ops.constant([2, 2], dtype=dtypes.int64))
|
|
|
|
return dataset.map(make_sparse)
|
|
|
|
dataset = iter(
|
|
strategy.distribute_datasets_from_function(
|
|
dataset_fn,
|
|
distribute_lib.InputOptions(experimental_fetch_to_device=False)))
|
|
|
|
sparse, result = sparse_lookup(dataset)
|
|
|
|
# All replicas return identical results.
|
|
for replica in range(strategy.num_replicas_in_sync):
|
|
self.assertIsInstance(sparse[replica], sparse_tensor.SparseTensor)
|
|
self.assertAllEqual(sparse[replica].indices, [[0, 0], [1, 0], [1, 1]])
|
|
self.assertAllEqual(sparse[replica].values, [0, 0, 1])
|
|
self.assertAllEqual(sparse[replica].dense_shape, [2, 2])
|
|
self.assertAllEqual(result[replica], [[0.0, 1.0], [3.0, 8.0]])
|
|
|
|
def test_composite_input_non_flat_output(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
if strategy.num_replicas_in_sync != 2:
|
|
self.skipTest("Test assumes two replicas.")
|
|
|
|
with strategy.scope():
|
|
table = variables.Variable(
|
|
initial_value=[[0.0, 1.0], [3.0, 7.0]], dtype=dtypes.float32)
|
|
|
|
@def_function.function
|
|
def sparse_lookup(iterator):
|
|
|
|
def tpu_function(sparse):
|
|
# Assumes dense_shape is (2, *)
|
|
looked_up = array_ops.gather(table, sparse.values)
|
|
segment_sum = math_ops.unsorted_segment_sum(
|
|
looked_up, sparse.indices[:, 0], 2)
|
|
return {"sparse": sparse, "segment_sum": segment_sum}
|
|
|
|
return nest.map_structure(
|
|
strategy.experimental_local_results,
|
|
strategy.run(tpu_function, args=(next(iterator),)))
|
|
|
|
def dataset_fn(_):
|
|
dataset = dataset_ops.Dataset.range(2)
|
|
|
|
def make_sparse(_):
|
|
return sparse_tensor.SparseTensor(
|
|
indices=array_ops.constant([[0, 0], [1, 0], [1, 1]],
|
|
dtype=dtypes.int64),
|
|
values=array_ops.constant([0, 0, 1], dtype=dtypes.int32),
|
|
dense_shape=array_ops.constant([2, 2], dtype=dtypes.int64))
|
|
|
|
return dataset.map(make_sparse)
|
|
|
|
dataset = iter(
|
|
strategy.distribute_datasets_from_function(
|
|
dataset_fn,
|
|
distribute_lib.InputOptions(experimental_fetch_to_device=False)))
|
|
|
|
output = sparse_lookup(dataset)
|
|
|
|
# All replicas return identical results.
|
|
for replica in range(strategy.num_replicas_in_sync):
|
|
self.assertIsInstance(output["sparse"][replica],
|
|
sparse_tensor.SparseTensor)
|
|
self.assertAllEqual(output["sparse"][replica].indices,
|
|
[[0, 0], [1, 0], [1, 1]])
|
|
self.assertAllEqual(output["sparse"][replica].values, [0, 0, 1])
|
|
self.assertAllEqual(output["sparse"][replica].dense_shape, [2, 2])
|
|
self.assertAllEqual(output["segment_sum"][replica],
|
|
[[0.0, 1.0], [3.0, 8.0]])
|
|
|
|
def test_composite_input_dynamic_shapes_outside_compilation(
|
|
self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
if strategy.num_replicas_in_sync != 2:
|
|
self.skipTest("Test assumes two replicas.")
|
|
|
|
table = variables.Variable(
|
|
initial_value=[[0.0, 1.0], [3.0, 7.0]], dtype=dtypes.float32)
|
|
|
|
@def_function.function
|
|
def sparse_lookup(iterator):
|
|
|
|
def tpu_function(sparse):
|
|
lookup = tpu_replication.outside_compilation(
|
|
embedding_ops.safe_embedding_lookup_sparse, table, sparse)
|
|
return math_ops.reduce_sum(lookup, axis=0)
|
|
|
|
return strategy.experimental_local_results(
|
|
strategy.run(tpu_function, args=(next(iterator),)))
|
|
|
|
def dataset_fn(_):
|
|
dataset = dataset_ops.Dataset.range(2)
|
|
|
|
def make_sparse(i):
|
|
indices = array_ops.constant([[0, 0], [1, 0], [1, 1]],
|
|
dtype=dtypes.int64)[0:2 + i]
|
|
values = array_ops.constant([0, 0, 1], dtype=dtypes.int32)[0:2 + i]
|
|
shape = [
|
|
array_ops.constant([2], dtype=dtypes.int64),
|
|
array_ops.expand_dims(1 + i, axis=0)
|
|
]
|
|
dense_shape = array_ops.concat(shape, axis=0)
|
|
return sparse_tensor.SparseTensor(
|
|
indices=indices, values=values, dense_shape=dense_shape)
|
|
|
|
return dataset.map(make_sparse)
|
|
|
|
dataset = iter(
|
|
strategy.distribute_datasets_from_function(
|
|
dataset_fn,
|
|
options=distribute_lib.InputOptions(
|
|
experimental_fetch_to_device=False)))
|
|
|
|
result = sparse_lookup(dataset)
|
|
self.assertAllEqual(result, [[0.0, 2.0], [1.5, 5.0]])
|
|
|
|
def test_composite_input_with_non_flat_components(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
|
|
class TestCompositeTypeSpec(type_spec.TypeSpec):
|
|
|
|
def __init__(self, component_type_spec):
|
|
self._component_type_spec = component_type_spec
|
|
|
|
@property
|
|
def value_type(self):
|
|
return TestComposite
|
|
|
|
def _to_components(self, value):
|
|
return value.values
|
|
|
|
def _from_components(self, components):
|
|
return TestComposite(components[0], components[1][0], components[1][1])
|
|
|
|
@property
|
|
def _component_specs(self):
|
|
return [self._component_type_spec,
|
|
[self._component_type_spec, self._component_type_spec]]
|
|
|
|
def _serialize(self):
|
|
return (self._component_type_spec,)
|
|
|
|
class TestComposite(composite_tensor.CompositeTensor):
|
|
|
|
def __init__(self, value1, value2, value3):
|
|
self.values = [value1, [value2, value3]]
|
|
|
|
@property
|
|
def _type_spec(self):
|
|
return TestCompositeTypeSpec(
|
|
tensor_spec.TensorSpec.from_tensor(self.values[0]))
|
|
|
|
def _shape_invariant_to_type_spec(self, shape):
|
|
return [shape, [shape, shape]]
|
|
|
|
@def_function.function
|
|
def test_fn(test_composite):
|
|
|
|
def tpu_function(composite):
|
|
return (composite,
|
|
composite.values[0] + (
|
|
composite.values[1][0] + composite.values[1][1])/2)
|
|
|
|
return nest.map_structure(
|
|
strategy.experimental_local_results,
|
|
strategy.run(tpu_function, args=(test_composite,)))
|
|
|
|
a = array_ops.constant([0.1])
|
|
b = array_ops.constant([1.2])
|
|
c = array_ops.constant([-0.4])
|
|
test_composite = TestComposite(a, b, c)
|
|
|
|
composite, result = test_fn(test_composite)
|
|
|
|
# All replicas return identical results.
|
|
for replica in range(strategy.num_replicas_in_sync):
|
|
self.assertIsInstance(composite[replica], TestComposite)
|
|
self.assertAllEqual(composite[replica].values[0], a)
|
|
self.assertAllEqual(composite[replica].values[1][0], b)
|
|
self.assertAllEqual(composite[replica].values[1][1], c)
|
|
self.assertAllEqual(result[replica], array_ops.constant([0.50000006]))
|
|
|
|
def test_per_device_tracing_of_mirrored_variables(self, enable_packed_var):
|
|
# Define trace_count as a list to avoid python scoping error
|
|
trace_count = [0]
|
|
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
with strategy.scope():
|
|
variable = variables.Variable(0.0)
|
|
|
|
@def_function.function
|
|
def add_one():
|
|
trace_count[0] = trace_count[0] + 1
|
|
return math_ops.add(variable, constant_op.constant(1.0))
|
|
|
|
@def_function.function
|
|
def update_variable():
|
|
for device in set(strategy.extended.worker_devices):
|
|
with ops.device(device):
|
|
add_one()
|
|
|
|
with strategy.scope():
|
|
update_variable.get_concrete_function()
|
|
self.assertLen(strategy.extended.worker_devices, trace_count[0])
|
|
|
|
def test_tpu_cancellation_does_not_close_chips(self, enable_packed_var):
|
|
if tpu_lib.enable_batch_variable_initialization():
|
|
self.skipTest("b/271767559")
|
|
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
num_replicas = strategy.num_replicas_in_sync
|
|
with strategy.scope():
|
|
x = random_ops.random_normal((10240, 10240))
|
|
y = random_ops.random_normal((10240, 10240))
|
|
|
|
v = variables.Variable(array_ops.identity(x))
|
|
dist_dataset = strategy.experimental_distribute_dataset(
|
|
dataset_ops.Dataset.from_tensors(y).repeat(num_replicas).batch(
|
|
num_replicas))
|
|
dist_iterator = iter(dist_dataset)
|
|
|
|
@def_function.function
|
|
def train_steps(v, iterator, steps):
|
|
|
|
def step_fn(inputs):
|
|
for val in inputs:
|
|
v.assign(math_ops.matmul(v, val))
|
|
|
|
for _ in math_ops.range(steps):
|
|
strategy.run(step_fn, args=(next(iterator),))
|
|
|
|
with self.assertRaises(errors.OutOfRangeError):
|
|
# The iterator has num_replicas/num_replicas = 1 step only.
|
|
train_steps(v, dist_iterator, 2)
|
|
|
|
# If TPU chips are not closed we can run the function on TPU again.
|
|
w = variables.Variable(array_ops.identity(x))
|
|
dist_dataset = strategy.experimental_distribute_dataset(
|
|
dataset_ops.Dataset.from_tensors(y).repeat(num_replicas).batch(
|
|
num_replicas))
|
|
dist_iterator = iter(dist_dataset)
|
|
train_steps(w, dist_iterator, 1)
|
|
|
|
def test_tpu_hardware_feature(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
self.assertIsInstance(
|
|
strategy.extended.tpu_hardware_feature.embedding_feature,
|
|
tpu_hardware_feature.HardwareFeature.EmbeddingFeature)
|
|
|
|
def test_get_tpu_cluster_resolver(self, enable_packed_var):
|
|
strategy = get_tpu_strategy(enable_packed_var)
|
|
self.assertIsNotNone(strategy.cluster_resolver)
|
|
|
|
def test_replica_order_for_distribute_datasets_from_function(
|
|
self, enable_packed_var
|
|
):
|
|
def _create_dataset(strategy):
|
|
def dataset_fn(ctx):
|
|
del ctx
|
|
return dataset_ops.Dataset.range(2)
|
|
|
|
return strategy.distribute_datasets_from_function(dataset_fn)
|
|
|
|
values = self._test_replica_order(_create_dataset).values
|
|
|
|
self.assertLen(values, 2)
|
|
self.assertEqual(1, values[0].numpy())
|
|
self.assertEqual(0, values[1].numpy())
|
|
|
|
def test_replica_order_for_experimental_distribute_dataset(
|
|
self, enable_packed_var
|
|
):
|
|
def _create_dataset(strategy):
|
|
dataset = dataset_ops.Dataset.range(2).batch(2)
|
|
return strategy.experimental_distribute_dataset(dataset)
|
|
|
|
values = self._test_replica_order(_create_dataset).values
|
|
|
|
self.assertLen(values, 2)
|
|
self.assertEqual(1, values[0].numpy())
|
|
self.assertEqual(0, values[1].numpy())
|
|
|
|
def _test_replica_order(self, create_dist_dataset_fn):
|
|
tf2.enable()
|
|
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
device_assignment = device_assignment_lib.DeviceAssignment(
|
|
topology, core_assignment=[[[0, 0, 0, 1]], [[0, 0, 0, 0]]]
|
|
)
|
|
strategy = tpu_lib.TPUStrategyV2(
|
|
resolver, experimental_device_assignment=device_assignment
|
|
)
|
|
strategy.extended._enable_data_reorder = True
|
|
|
|
dist_dataset = create_dist_dataset_fn(strategy)
|
|
iterator = iter(dist_dataset)
|
|
|
|
@def_function.function
|
|
def test_iterators_order(iterator):
|
|
return next(iterator)
|
|
|
|
return test_iterators_order(iterator)
|
|
|
|
|
|
@test_util.with_eager_op_as_function
|
|
class TPUStrategyDataPrefetchTest(test.TestCase):
|
|
|
|
def test_prefetch_to_device_default(self):
|
|
strategy = get_tpu_strategy()
|
|
dataset = dataset_ops.Dataset.range(
|
|
strategy.num_replicas_in_sync * 2,
|
|
output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
|
|
|
|
# Check default, should prefetch to TPU.
|
|
dataset_item = next(iter(strategy.experimental_distribute_dataset(dataset)))
|
|
dataset_location = tf_device.DeviceSpec.from_string(
|
|
dataset_item.values[0].device)
|
|
self.assertEqual(dataset_location.device_type, "TPU")
|
|
|
|
def test_prefetch_to_device_tpu(self):
|
|
strategy = get_tpu_strategy()
|
|
dataset = dataset_ops.Dataset.range(
|
|
strategy.num_replicas_in_sync * 2,
|
|
output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
|
|
|
|
input_options = distribute_lib.InputOptions(
|
|
experimental_fetch_to_device=True)
|
|
dataset_item = next(iter(strategy.experimental_distribute_dataset(
|
|
dataset, options=input_options)))
|
|
dataset_location = tf_device.DeviceSpec.from_string(
|
|
dataset_item.values[0].device)
|
|
self.assertEqual(dataset_location.device_type, "TPU")
|
|
|
|
def test_prefetch_to_device_cpu(self):
|
|
strategy = get_tpu_strategy()
|
|
dataset = dataset_ops.Dataset.range(
|
|
strategy.num_replicas_in_sync * 2,
|
|
output_type=dtypes.float32).batch(strategy.num_replicas_in_sync)
|
|
|
|
# Should be CPU when prefetch_to_device is False.
|
|
input_options = distribute_lib.InputOptions(
|
|
experimental_fetch_to_device=False)
|
|
dataset_item = next(iter(strategy.experimental_distribute_dataset(
|
|
dataset, options=input_options)))
|
|
dataset_location = tf_device.DeviceSpec.from_string(
|
|
dataset_item.values[0].device)
|
|
self.assertEqual(dataset_location.device_type, "CPU")
|
|
|
|
def test_prefetch_to_device_sparse_dataset(self):
|
|
strategy = get_tpu_strategy()
|
|
# Values here aren't important.
|
|
dataset = dataset_ops.Dataset.from_tensors(
|
|
sparse_tensor.SparseTensor(indices=[[0, 0], [0, 1], [1, 0]],
|
|
values=[1, 2, 3],
|
|
dense_shape=[2, 2]))
|
|
dataset = dataset.repeat()
|
|
dataset = dataset.batch(strategy.num_replicas_in_sync)
|
|
|
|
with self.assertRaisesRegex(ValueError, "TPUStrategy does not support"):
|
|
iter(strategy.experimental_distribute_dataset(dataset))
|
|
|
|
def test_prefetch_to_device_ragged_dataset(self):
|
|
strategy = get_tpu_strategy()
|
|
# Values here aren't important.
|
|
dataset = dataset_ops.Dataset.from_tensors(
|
|
ragged_tensor.RaggedTensor.from_row_splits(
|
|
values=[1, 2, 3],
|
|
row_splits=[0, 2, 3]))
|
|
dataset = dataset.repeat()
|
|
dataset = dataset.batch(strategy.num_replicas_in_sync)
|
|
|
|
with self.assertRaisesRegex(ValueError, "TPUStrategy does not support"):
|
|
iter(strategy.experimental_distribute_dataset(dataset))
|
|
|
|
def test_prefetch_to_device_sparse_dataset_fn(self):
|
|
strategy = get_tpu_strategy()
|
|
def dataset_fn(ctx):
|
|
del ctx
|
|
# Values here aren't important.
|
|
dataset = dataset_ops.Dataset.from_tensors(
|
|
sparse_tensor.SparseTensor(indices=[[0, 0], [0, 1], [1, 0]],
|
|
values=[1, 2, 3],
|
|
dense_shape=[2, 2]))
|
|
dataset = dataset.repeat()
|
|
return dataset.batch(strategy.num_replicas_in_sync)
|
|
|
|
with self.assertRaisesRegex(ValueError, "TPUStrategy does not support"):
|
|
iter(strategy.distribute_datasets_from_function(dataset_fn))
|
|
|
|
def test_prefetch_to_device_ragged_dataset_fn(self):
|
|
strategy = get_tpu_strategy()
|
|
def dataset_fn(ctx):
|
|
del ctx
|
|
# Values here aren't important.
|
|
dataset = dataset_ops.Dataset.from_tensors(
|
|
ragged_tensor.RaggedTensor.from_row_splits(
|
|
values=[1, 2, 3],
|
|
row_splits=[0, 2, 3]))
|
|
dataset = dataset.repeat()
|
|
return dataset.batch(strategy.num_replicas_in_sync)
|
|
|
|
with self.assertRaisesRegex(ValueError, "TPUStrategy does not support"):
|
|
iter(strategy.distribute_datasets_from_function(dataset_fn))
|
|
|
|
def test_create_iterator_on_device(self):
|
|
|
|
@def_function.function
|
|
def create_iter():
|
|
with ops.device("/device:TPU:0"):
|
|
return gen_dataset_ops.anonymous_iterator_v3(
|
|
output_types=[dtypes.float32], output_shapes=[[]])
|
|
|
|
create_iter()
|
|
|
|
|
|
@test_util.with_eager_op_as_function
|
|
class TPUStrategyDistributionTest(
|
|
strategy_test_lib.DistributionTestBase,
|
|
strategy_test_lib.TwoDeviceDistributionTestBase):
|
|
|
|
def test_update_config_proto(self):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
strategy = tpu_lib.TPUStrategyV2(resolver)
|
|
|
|
config_proto = config_pb2.ConfigProto()
|
|
cluster_spec = server_lib.ClusterSpec({"worker": ["fake1", "fake2"]})
|
|
with test.mock.patch.object(
|
|
resolver, "cluster_spec", return_value=cluster_spec):
|
|
new_config = strategy.update_config_proto(config_proto)
|
|
|
|
# Verify cluster_def.
|
|
self.assertProtoEquals(cluster_spec.as_cluster_def(),
|
|
new_config.cluster_def)
|
|
|
|
# Verify isolate_session_state
|
|
self.assertTrue(new_config.isolate_session_state)
|
|
|
|
def test_make_input_fn_iterable(self):
|
|
dataset_fn = lambda: dataset_ops.Dataset.range(10)
|
|
expected_values = [[i, i+1] for i in range(0, 10, 2)]
|
|
distribution = get_tpu_strategy()
|
|
input_fn = self._input_fn_to_test_input_context(
|
|
dataset_fn,
|
|
expected_num_replicas_in_sync=2,
|
|
expected_num_input_pipelines=1,
|
|
expected_input_pipeline_id=0)
|
|
self._test_input_fn_iterable(distribution, input_fn, expected_values)
|
|
|
|
def test_make_input_fn_iterator(self):
|
|
dataset_fn = lambda: dataset_ops.Dataset.range(10)
|
|
expected_values = [[i, i+1] for i in range(0, 10, 2)]
|
|
distribution = get_tpu_strategy()
|
|
input_fn = self._input_fn_to_test_input_context(
|
|
dataset_fn,
|
|
expected_num_replicas_in_sync=2,
|
|
expected_num_input_pipelines=1,
|
|
expected_input_pipeline_id=0)
|
|
iterator = distribution.make_input_fn_iterator(input_fn)
|
|
self._test_input_fn_iterator(
|
|
iterator,
|
|
distribution.extended.worker_devices,
|
|
expected_values)
|
|
|
|
def test_num_replicas_in_sync(self):
|
|
strategy = get_tpu_strategy()
|
|
self.assertEqual(2, strategy.num_replicas_in_sync)
|
|
|
|
def test_call_and_merge_exceptions(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_call_and_merge_exceptions(strategy)
|
|
|
|
def test_numpy_dataset(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_numpy_dataset(strategy, run_in_function=True)
|
|
|
|
def test_global_step_update(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_global_step_update(strategy)
|
|
|
|
def test_run(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_run(strategy, run_in_function=True)
|
|
|
|
def test_summary_for_replica_zero_only(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_summary_for_replica_zero_only(strategy)
|
|
|
|
def test_all_reduce_sum(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_all_reduce_sum(strategy, run_in_function=True)
|
|
|
|
def test_all_reduce_sum_gradients(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_all_reduce_sum_gradients(strategy, run_in_function=True)
|
|
|
|
def test_all_reduce_sum_gradient_tape(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_all_reduce_sum_gradient_tape(strategy, run_in_function=True)
|
|
|
|
def test_all_reduce_mean(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_all_reduce_mean(strategy, run_in_function=True)
|
|
|
|
def test_all_reduce_mean_gradients(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_all_reduce_mean_gradients(strategy, run_in_function=True)
|
|
|
|
def test_all_reduce_mean_gradient_tape(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_all_reduce_mean_gradient_tape(strategy, run_in_function=True)
|
|
|
|
def test_reduce(self):
|
|
strategy = get_tpu_strategy()
|
|
|
|
inputs = strategy.make_input_fn_iterator(
|
|
lambda _: dataset_ops.Dataset.from_tensor_slices([2., 3.]))
|
|
|
|
self.evaluate(inputs.initialize())
|
|
per_replica_outputs = strategy.run(
|
|
def_function.function(math_ops.square), args=(next(inputs),))
|
|
|
|
with strategy.scope():
|
|
mean = strategy.reduce(reduce_util.ReduceOp.MEAN, per_replica_outputs,
|
|
axis=None)
|
|
self.assertEqual(6.5, self.evaluate(mean))
|
|
|
|
def test_constraint(self):
|
|
strategy = get_tpu_strategy()
|
|
|
|
with strategy.scope():
|
|
variable = variables.Variable(initial_value=2.,
|
|
constraint=lambda x: 0. * x + 1.)
|
|
self.assertEqual(variable.value().numpy(), 2)
|
|
|
|
@def_function.function
|
|
def update_variable():
|
|
variable.assign_add(1)
|
|
variable.assign(variable.constraint(variable))
|
|
|
|
update_variable()
|
|
self.assertEqual(variable.value().numpy(), 1)
|
|
|
|
def test_trainable_variables(self):
|
|
strategy = get_tpu_strategy()
|
|
self._test_trainable_variable(strategy)
|
|
|
|
|
|
@test_util.with_eager_op_as_function
|
|
class DeviceAssignmentTest(test.TestCase):
|
|
|
|
def test_core_assignment(self):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
device_assignment = device_assignment_lib.DeviceAssignment(
|
|
topology, core_assignment=[[[0, 0, 0, 0]]])
|
|
self.assertAllEqual([[[0, 0, 0, 0]]], device_assignment.core_assignment)
|
|
self.assertEqual(1, device_assignment.num_cores_per_replica)
|
|
self.assertEqual(1, device_assignment.num_replicas)
|
|
self.assertEqual("/task:0/device:TPU:0", device_assignment.tpu_device())
|
|
self.assertEqual("/task:0/device:CPU:0", device_assignment.host_device())
|
|
|
|
def test_device_assignment_strategy_properties(self):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
device_assignment = device_assignment_lib.DeviceAssignment(
|
|
topology, core_assignment=[[[0, 0, 0, 0]]])
|
|
strategy = tpu_lib.TPUStrategyV2(
|
|
resolver,
|
|
experimental_device_assignment=device_assignment)
|
|
self.assertEqual(strategy.extended.num_hosts, 1)
|
|
self.assertEqual(strategy.num_replicas_in_sync, 1)
|
|
self.assertEqual(strategy.extended.num_replicas_per_host, 1) # pylint: disable=protected-access
|
|
|
|
def test_device_assignment_constants(self):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
device_assignment = device_assignment_lib.DeviceAssignment(
|
|
topology,
|
|
core_assignment=device_assignment_lib.SINGLE_CORE_ASSIGNMENT)
|
|
self.assertAllEqual([[[0, 0, 0, 0]]], device_assignment.core_assignment)
|
|
self.assertEqual(1, device_assignment.num_cores_per_replica)
|
|
self.assertEqual(1, device_assignment.num_replicas)
|
|
self.assertEqual("/task:0/device:TPU:0", device_assignment.tpu_device())
|
|
self.assertEqual("/task:0/device:CPU:0", device_assignment.host_device())
|
|
|
|
def test_variables_mismatched_device_assignment(self):
|
|
resolver = get_tpu_cluster_resolver()
|
|
remote.connect_to_cluster(resolver)
|
|
topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
|
|
|
|
strategy0 = tpu_lib.TPUStrategyV2(resolver)
|
|
self.assertEqual(
|
|
("/job:localhost/replica:0/task:0/device:TPU:0",
|
|
"/job:localhost/replica:0/task:0/device:TPU:1"),
|
|
strategy0.extended.worker_devices)
|
|
|
|
with strategy0.scope():
|
|
v = variables.Variable(1.)
|
|
|
|
v1_assign_op = strategy0.experimental_local_results(v)[1].assign(42.)
|
|
|
|
with self.cached_session():
|
|
self.evaluate(variables.global_variables_initializer())
|
|
self.evaluate(v1_assign_op)
|
|
self.assertAllEqual([1., 42.],
|
|
self.evaluate(
|
|
strategy0.experimental_local_results(v)))
|
|
|
|
# Second strategy has devices reversed relative to the first.
|
|
device_assignment = device_assignment_lib.DeviceAssignment(
|
|
topology, core_assignment=[[[0, 0, 0, 1]], [[0, 0, 0, 0]]])
|
|
strategy1 = tpu_lib.TPUStrategyV2(
|
|
resolver,
|
|
experimental_device_assignment=device_assignment)
|
|
self.assertEqual(
|
|
("/job:localhost/replica:0/task:0/device:TPU:1",
|
|
"/job:localhost/replica:0/task:0/device:TPU:0"),
|
|
strategy1.extended.worker_devices)
|
|
|
|
v_read = strategy1.run(def_function.function(v.read_value))
|
|
|
|
with self.cached_session():
|
|
self.assertAllEqual([42., 1.],
|
|
self.evaluate(
|
|
strategy0.experimental_local_results(v_read)))
|
|
|
|
|
|
class VariableCreationTest(test.TestCase):
|
|
|
|
def test_custom_tpu_variable_creator(self):
|
|
strategy = get_tpu_strategy()
|
|
|
|
def variable_creator(next_creator, **kwargs):
|
|
def custom_tpu_variable_creator(next_creator, **kwargs):
|
|
return next_creator(**kwargs)
|
|
kwargs["custom_tpu_variable_creator"] = custom_tpu_variable_creator
|
|
return next_creator(**kwargs)
|
|
|
|
with strategy.scope():
|
|
tpu_variable = variables.Variable(1.0)
|
|
|
|
self.assertIsInstance(tpu_variable, tpu_values.TPUDistributedVariable)
|
|
|
|
with variable_scope.variable_creator_scope(variable_creator):
|
|
non_tpu_variable = variables.Variable(1.0)
|
|
self.assertIsInstance(non_tpu_variable, variables.Variable)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ops.enable_eager_execution()
|
|
test.main()
|