wehub-resource-sync
913 lines
34 KiB
Python
913 lines
34 KiB
Python
# Copyright 2019 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 that the system configuration methods work properly."""
|
|
|
|
from absl.testing import parameterized
|
|
|
|
from tensorflow.core.protobuf import cluster_pb2
|
|
from tensorflow.core.protobuf import config_pb2
|
|
from tensorflow.core.protobuf import rewriter_config_pb2
|
|
from tensorflow.core.protobuf import tensorflow_server_pb2
|
|
from tensorflow.python.eager import context
|
|
from tensorflow.python.eager import def_function
|
|
from tensorflow.python.framework import config
|
|
from tensorflow.python.framework import constant_op
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import errors
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import test_ops
|
|
from tensorflow.python.framework import test_util
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.platform import test
|
|
from tensorflow.python.util import compat
|
|
|
|
|
|
def reset_eager(fn):
|
|
|
|
def wrapper(*args, **kwargs):
|
|
try:
|
|
return fn(*args, **kwargs)
|
|
finally:
|
|
# Reset the context.
|
|
context._reset_jit_compiler_flags()
|
|
context._reset_context()
|
|
ops.enable_eager_execution_internal()
|
|
assert context._context is not None
|
|
|
|
return wrapper
|
|
|
|
|
|
@test_util.with_eager_op_as_function
|
|
class ConfigTest(test.TestCase, parameterized.TestCase):
|
|
|
|
@test_util.disable_eager_op_as_function('b/204320409')
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testDevicePolicy(self):
|
|
self.assertEqual(context.DEVICE_PLACEMENT_SILENT,
|
|
context.context().device_policy)
|
|
|
|
# If no op has been executed we should be able to set the device policy as
|
|
# well as any init-time configs.
|
|
config.set_intra_op_parallelism_threads(1)
|
|
config.set_device_policy('silent')
|
|
config.set_intra_op_parallelism_threads(2)
|
|
|
|
context.ensure_initialized()
|
|
|
|
def copy_tensor(dtype=dtypes.int32):
|
|
with ops.device('CPU:0'):
|
|
cpu_tensor = constant_op.constant(1, dtype=dtype)
|
|
gpu_tensor = cpu_tensor.gpu()
|
|
self.assertAllEqual(cpu_tensor + gpu_tensor, 2.0)
|
|
|
|
config.set_device_policy('silent')
|
|
self.assertEqual(config.get_device_policy(), 'silent')
|
|
self.assertEqual(context.DEVICE_PLACEMENT_SILENT,
|
|
context.context().device_policy)
|
|
copy_tensor()
|
|
|
|
config.set_device_policy('silent_for_int32')
|
|
self.assertEqual(config.get_device_policy(), 'silent_for_int32')
|
|
self.assertEqual(context.DEVICE_PLACEMENT_SILENT_FOR_INT32,
|
|
context.context().device_policy)
|
|
with self.assertRaisesRegex(errors.InvalidArgumentError,
|
|
'Tensors on conflicting devices'):
|
|
copy_tensor(dtypes.float32)
|
|
copy_tensor()
|
|
|
|
config.set_device_policy('warn')
|
|
self.assertEqual(config.get_device_policy(), 'warn')
|
|
self.assertEqual(context.DEVICE_PLACEMENT_WARN,
|
|
context.context().device_policy)
|
|
copy_tensor()
|
|
|
|
config.set_device_policy('explicit')
|
|
self.assertEqual(config.get_device_policy(), 'explicit')
|
|
self.assertEqual(context.DEVICE_PLACEMENT_EXPLICIT,
|
|
context.context().device_policy)
|
|
with self.assertRaisesRegex(errors.InvalidArgumentError,
|
|
'Tensors on conflicting devices'):
|
|
copy_tensor()
|
|
|
|
config.set_device_policy(None)
|
|
self.assertEqual(config.get_device_policy(), 'silent')
|
|
|
|
@reset_eager
|
|
def testExecutionMode(self):
|
|
self.assertTrue(config.get_synchronous_execution())
|
|
self.assertEqual(context.SYNC, context.context().execution_mode)
|
|
|
|
# If no op has been executed we should be able to set the execution mode as
|
|
# well as any init-time configs.
|
|
config.set_intra_op_parallelism_threads(1)
|
|
config.set_synchronous_execution(False)
|
|
config.set_intra_op_parallelism_threads(2)
|
|
|
|
config.set_synchronous_execution(True)
|
|
self.assertTrue(config.get_synchronous_execution())
|
|
self.assertEqual(context.SYNC, context.context().execution_mode)
|
|
config.set_synchronous_execution(False)
|
|
self.assertFalse(config.get_synchronous_execution())
|
|
self.assertEqual(context.ASYNC, context.context().execution_mode)
|
|
|
|
@reset_eager
|
|
def testIntraOpParallelismThreads(self):
|
|
config.set_intra_op_parallelism_threads(10)
|
|
self.assertEqual(config.get_intra_op_parallelism_threads(),
|
|
context.context().intra_op_parallelism_threads)
|
|
|
|
context.ensure_initialized()
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
config.set_intra_op_parallelism_threads(1)
|
|
|
|
config.set_intra_op_parallelism_threads(10)
|
|
|
|
@reset_eager
|
|
def testInterOpParallelismThreads(self):
|
|
config.set_inter_op_parallelism_threads(10)
|
|
self.assertEqual(config.get_inter_op_parallelism_threads(),
|
|
context.context().inter_op_parallelism_threads)
|
|
|
|
context.ensure_initialized()
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
config.set_inter_op_parallelism_threads(1)
|
|
|
|
config.set_inter_op_parallelism_threads(10)
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testSoftPlacement(self):
|
|
if context.executing_eagerly():
|
|
self.assertTrue(config.get_soft_device_placement())
|
|
else:
|
|
self.assertFalse(config.get_soft_device_placement())
|
|
|
|
def test_attr():
|
|
with ops.device('/device:GPU:0'):
|
|
return test_ops.test_attr(T=dtypes.float32, name='test_attr')
|
|
|
|
config.set_soft_device_placement(True)
|
|
self.assertEqual(config.get_soft_device_placement(), True)
|
|
self.assertEqual(config.get_soft_device_placement(),
|
|
context.context().soft_device_placement)
|
|
|
|
# Since soft placement is enabled, the test_attr operation should fallback
|
|
# to CPU with pure eager execution as well as functions
|
|
test_attr()
|
|
def_function.function(test_attr)()
|
|
|
|
config.set_soft_device_placement(False)
|
|
self.assertEqual(config.get_soft_device_placement(), False)
|
|
self.assertEqual(config.get_soft_device_placement(),
|
|
context.context().soft_device_placement)
|
|
|
|
# Since soft placement is disabled, the test_attr operation should fail on
|
|
# GPU with pure eager execution as well as functions
|
|
with self.assertRaises(errors.InvalidArgumentError):
|
|
test_attr()
|
|
with self.assertRaises(errors.InvalidArgumentError):
|
|
def_function.function(test_attr)()
|
|
|
|
@reset_eager
|
|
def testLogDevicePlacement(self):
|
|
self.assertFalse(context.get_log_device_placement())
|
|
|
|
context.set_log_device_placement(True)
|
|
self.assertEqual(context.get_log_device_placement(), True)
|
|
self.assertEqual(context.get_log_device_placement(),
|
|
context.context().log_device_placement)
|
|
|
|
context.set_log_device_placement(False)
|
|
self.assertEqual(context.get_log_device_placement(), False)
|
|
self.assertEqual(context.get_log_device_placement(),
|
|
context.context().log_device_placement)
|
|
|
|
context.ensure_initialized()
|
|
|
|
# Changing the device placement should not throw an exception
|
|
context.set_log_device_placement(True)
|
|
|
|
@reset_eager
|
|
def testEnableMlirBridge(self):
|
|
# Default value of enable_mlir_bridge is false.
|
|
self.assertFalse(context.context().config.experimental.enable_mlir_bridge)
|
|
self.assertEqual(
|
|
context.context().config.experimental.mlir_bridge_rollout,
|
|
config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_UNSPECIFIED)
|
|
|
|
# Tests enabling mlir bridge.
|
|
config.enable_mlir_bridge()
|
|
self.assertTrue(context.context().config.experimental.enable_mlir_bridge)
|
|
self.assertEqual(
|
|
context.context().config.experimental.mlir_bridge_rollout,
|
|
config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_ENABLED)
|
|
|
|
# Tests disabling mlir bridge.
|
|
config.disable_mlir_bridge()
|
|
self.assertFalse(context.context().config.experimental.enable_mlir_bridge)
|
|
self.assertEqual(
|
|
context.context().config.experimental.mlir_bridge_rollout,
|
|
config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_DISABLED)
|
|
|
|
@reset_eager
|
|
def testResetMlirFlags(self):
|
|
# Default value of enable_mlir_bridge is false.
|
|
self.assertFalse(context.context().config.experimental.enable_mlir_bridge)
|
|
self.assertEqual(
|
|
context.context().config.experimental.mlir_bridge_rollout,
|
|
config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_UNSPECIFIED)
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testJit(self):
|
|
self.assertEqual(config.get_optimizer_jit(), '')
|
|
|
|
# the following function should cause Op fusion to occur. However, there is
|
|
# unfortunately no straightforward way to ensure this. We will just have to
|
|
# settle for creating a test that can trigger JIT.
|
|
@def_function.function
|
|
def fun(a, b):
|
|
c = a * b
|
|
d = c + a
|
|
return d
|
|
|
|
a = constant_op.constant([2., 2.])
|
|
b = constant_op.constant([2., 2.])
|
|
|
|
self.evaluate(fun(a, b))
|
|
|
|
config.set_optimizer_jit('autoclustering')
|
|
self.assertEqual(config.get_optimizer_jit(), 'autoclustering')
|
|
|
|
self.evaluate(fun(a, b))
|
|
|
|
config.set_optimizer_jit('')
|
|
self.assertEqual(config.get_optimizer_jit(), '')
|
|
|
|
self.evaluate(fun(a, b))
|
|
|
|
@parameterized.named_parameters(
|
|
('LayoutOptimizer', 'layout_optimizer'),
|
|
('ConstantFolding', 'constant_folding'),
|
|
('ShapeOptimization', 'shape_optimization'), ('Remapping', 'remapping'),
|
|
('ArithmeticOptimization', 'arithmetic_optimization'),
|
|
('DependencyOptimization', 'dependency_optimization'),
|
|
('LoopOptimization', 'loop_optimization'),
|
|
('FunctionOptimization', 'function_optimization'),
|
|
('DebugStripper', 'debug_stripper'),
|
|
('ScopedAllocatorOptimization', 'scoped_allocator_optimization'),
|
|
('ImplementationSelector', 'implementation_selector'),
|
|
('AutoMixedPrecision', 'auto_mixed_precision'))
|
|
@reset_eager
|
|
def testOptimizerToggleOption(self, field):
|
|
# TODO(b/128531235): Improve testing of option
|
|
options = config.get_optimizer_experimental_options()
|
|
self.assertIsNone(options.get(field))
|
|
|
|
config.set_optimizer_experimental_options({field: True})
|
|
options[field] = True
|
|
self.assertDictEqual(config.get_optimizer_experimental_options(), options)
|
|
self.assertDictEqual(context.context().get_optimizer_experimental_options(),
|
|
options)
|
|
|
|
config.set_optimizer_experimental_options({field: False})
|
|
options[field] = False
|
|
self.assertDictEqual(config.get_optimizer_experimental_options(), options)
|
|
self.assertDictEqual(context.context().get_optimizer_experimental_options(),
|
|
options)
|
|
|
|
@parameterized.named_parameters(
|
|
('DisableModelPruning', 'disable_model_pruning'),
|
|
('DisableMetaOptimizer', 'disable_meta_optimizer'))
|
|
@reset_eager
|
|
def testOptimizerBoolOption(self, field):
|
|
# TODO(b/128531235): Improve testing of option
|
|
options = config.get_optimizer_experimental_options()
|
|
self.assertFalse(options.get(field))
|
|
|
|
config.set_optimizer_experimental_options({field: True})
|
|
options[field] = True
|
|
self.assertDictEqual(config.get_optimizer_experimental_options(), options)
|
|
self.assertDictEqual(context.context().get_optimizer_experimental_options(),
|
|
options)
|
|
|
|
config.set_optimizer_experimental_options({field: False})
|
|
options[field] = False
|
|
self.assertDictEqual(config.get_optimizer_experimental_options(), options)
|
|
self.assertDictEqual(context.context().get_optimizer_experimental_options(),
|
|
options)
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testOptimizerToggleOptionPinToHost(self):
|
|
options = config.get_optimizer_experimental_options()
|
|
self.assertIsNone(options.get('pin_to_host_optimization'))
|
|
|
|
@def_function.function
|
|
def fun():
|
|
op = test_ops.device_placement_op()
|
|
return op
|
|
|
|
# Force optimizer to run for all graphs
|
|
config.set_optimizer_experimental_options({'min_graph_nodes': -1})
|
|
options['min_graph_nodes'] = -1
|
|
|
|
# Since pin to host is disabled, the operation should go on GPU
|
|
gpu = self.evaluate(fun())
|
|
self.assertIn(compat.as_bytes('GPU'), gpu)
|
|
|
|
config.set_optimizer_experimental_options(
|
|
{'pin_to_host_optimization': True})
|
|
options['pin_to_host_optimization'] = True
|
|
self.assertDictEqual(config.get_optimizer_experimental_options(), options)
|
|
self.assertDictEqual(context.context().get_optimizer_experimental_options(),
|
|
options)
|
|
|
|
# Since pin to host is enabled, the operation should go on CPU
|
|
cpu = self.evaluate(fun())
|
|
self.assertIn(compat.as_bytes('CPU'), cpu)
|
|
|
|
config.set_optimizer_experimental_options(
|
|
{'pin_to_host_optimization': False})
|
|
options['pin_to_host_optimization'] = False
|
|
self.assertDictEqual(config.get_optimizer_experimental_options(), options)
|
|
self.assertDictEqual(context.context().get_optimizer_experimental_options(),
|
|
options)
|
|
|
|
# Since pin to host is disabled again, the operation should go on GPU
|
|
gpu2 = self.evaluate(fun())
|
|
self.assertIn(compat.as_bytes('GPU'), gpu2)
|
|
|
|
|
|
class DeviceTest(test.TestCase):
|
|
|
|
@reset_eager
|
|
def testPhysicalDevices(self):
|
|
cpus = config.list_physical_devices('CPU')
|
|
self.assertGreater(len(cpus), 0)
|
|
if test_util.is_gpu_available():
|
|
gpus = config.list_physical_devices('GPU')
|
|
self.assertGreater(len(gpus), 0)
|
|
|
|
@reset_eager
|
|
def testCpuMultiple(self):
|
|
cpus = config.list_physical_devices('CPU')
|
|
self.assertEqual(len(cpus), 1)
|
|
|
|
config.set_logical_device_configuration(cpus[0], [
|
|
context.LogicalDeviceConfiguration(),
|
|
context.LogicalDeviceConfiguration()
|
|
])
|
|
|
|
context.ensure_initialized()
|
|
|
|
vcpus = config.list_logical_devices('CPU')
|
|
self.assertEqual(len(vcpus), 2)
|
|
|
|
with ops.device('/device:CPU:0'):
|
|
a = constant_op.constant(1.0)
|
|
self.evaluate(a)
|
|
with ops.device('/device:CPU:1'):
|
|
b = constant_op.constant(1.0)
|
|
self.evaluate(b)
|
|
with ops.device('/device:CPU:2'):
|
|
c = constant_op.constant(1.0)
|
|
self.evaluate(c)
|
|
if test_util.is_gpu_available():
|
|
self.assertIn('GPU:0', c.device)
|
|
else:
|
|
self.assertIn('CPU:0', c.device)
|
|
|
|
# Ensure we can place ops on each of the device names
|
|
for vcpu in vcpus:
|
|
with ops.device(vcpu.name):
|
|
d = constant_op.constant(1.0)
|
|
self.evaluate(d)
|
|
|
|
# Modifying the CPU configuration is not supported
|
|
with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
|
|
config.set_logical_device_configuration(cpus[0], [
|
|
context.LogicalDeviceConfiguration(),
|
|
context.LogicalDeviceConfiguration(),
|
|
context.LogicalDeviceConfiguration()
|
|
])
|
|
|
|
# Setting the same CPU configuration is fine
|
|
config.set_logical_device_configuration(cpus[0], [
|
|
context.LogicalDeviceConfiguration(),
|
|
context.LogicalDeviceConfiguration()
|
|
])
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testGpuNone(self):
|
|
config.set_soft_device_placement(False)
|
|
gpus = config.list_physical_devices('GPU')
|
|
self.assertGreater(len(gpus), 0)
|
|
|
|
cpus = config.list_physical_devices('CPU')
|
|
self.assertEqual(len(cpus), 1)
|
|
|
|
self.assertEqual(len(config.get_visible_devices('CPU')), 1)
|
|
self.assertGreater(len(config.get_visible_devices('GPU')), 0)
|
|
|
|
self.assertEqual(len(config.get_visible_devices('XLA_GPU')), 0)
|
|
|
|
config.set_visible_devices(cpus[0])
|
|
self.assertEqual(len(config.get_visible_devices('CPU')), 1)
|
|
self.assertEqual(len(config.get_visible_devices('GPU')), 0)
|
|
self.assertEqual(len(config.list_logical_devices('XLA_GPU')), 0)
|
|
|
|
with self.assertRaisesRegex(errors.InvalidArgumentError,
|
|
'Could not satisfy'):
|
|
with ops.device('/device:GPU:0'):
|
|
a = array_ops.identity(1.0)
|
|
self.evaluate(a)
|
|
|
|
# Modifying the visible devices is not supported
|
|
with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
|
|
config.set_visible_devices(gpus)
|
|
|
|
# Setting the same visible devices is fine
|
|
config.set_visible_devices(cpus[0])
|
|
|
|
@reset_eager
|
|
def testGpuMultiple(self):
|
|
config.set_soft_device_placement(False)
|
|
gpus = config.list_physical_devices('GPU')
|
|
|
|
if len(gpus) < 2:
|
|
self.skipTest('Need at least 2 GPUs')
|
|
|
|
context.ensure_initialized()
|
|
|
|
for i in range(0, len(gpus)):
|
|
with ops.device('/device:GPU:' + str(i)):
|
|
a = constant_op.constant(1.0)
|
|
self.evaluate(a)
|
|
|
|
with self.assertRaisesRegex(errors.InvalidArgumentError,
|
|
'Could not satisfy device specification'):
|
|
with ops.device('/device:GPU:' + str(len(gpus))):
|
|
a = constant_op.constant(1.0)
|
|
self.evaluate(a)
|
|
|
|
@reset_eager
|
|
def testDeviceDetails(self):
|
|
(cpu,) = config.list_physical_devices('CPU')
|
|
details = config.get_device_details(cpu)
|
|
self.assertEqual(details, {})
|
|
|
|
if not test_util.is_gpu_available():
|
|
return
|
|
|
|
gpus = config.list_physical_devices('GPU')
|
|
details = config.get_device_details(gpus[0])
|
|
self.assertIsInstance(details['device_name'], str)
|
|
self.assertNotEmpty(details['device_name'])
|
|
if test.is_built_with_rocm():
|
|
# AMD GPUs do not have a compute capability
|
|
self.assertNotIn('compute_capability', details)
|
|
else:
|
|
cc = details['compute_capability']
|
|
self.assertIsInstance(cc, tuple)
|
|
major, minor = cc
|
|
self.assertGreater(major, 0)
|
|
self.assertGreaterEqual(minor, 0)
|
|
|
|
# Test GPU returned from get_visible_devices
|
|
if len(gpus) > 2:
|
|
config.set_visible_devices(gpus[1], 'GPU')
|
|
(visible_gpu,) = config.get_visible_devices('GPU')
|
|
details = config.get_device_details(visible_gpu)
|
|
self.assertIsInstance(details['device_name'], str)
|
|
|
|
@reset_eager
|
|
def testDeviceDetailsErrors(self):
|
|
logical_devices = config.list_logical_devices()
|
|
with self.assertRaisesRegex(ValueError,
|
|
'must be a tf.config.PhysicalDevice'):
|
|
config.get_device_details(logical_devices[0])
|
|
|
|
phys_dev = context.PhysicalDevice('/physical_device:CPU:100', 'CPU')
|
|
with self.assertRaisesRegex(
|
|
ValueError, 'The PhysicalDevice must be one obtained from '
|
|
'calling `tf.config.list_physical_devices`'):
|
|
config.get_device_details(phys_dev)
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testVirtualGpu(self):
|
|
config.set_soft_device_placement(False)
|
|
gpus = config.list_physical_devices('GPU')
|
|
self.assertNotEqual(len(gpus), 0)
|
|
|
|
self.assertIsNone(config.get_logical_device_configuration(gpus[-1]))
|
|
config.set_logical_device_configuration(gpus[-1], [
|
|
context.LogicalDeviceConfiguration(memory_limit=10),
|
|
context.LogicalDeviceConfiguration(memory_limit=10)
|
|
])
|
|
self.assertEqual(len(config.get_logical_device_configuration(gpus[-1])), 2)
|
|
|
|
logical_gpus = config.list_logical_devices('GPU')
|
|
self.assertTrue(len(logical_gpus), len(gpus) + 1)
|
|
for i in range(0, len(logical_gpus)):
|
|
with ops.device('/device:GPU:' + str(i)):
|
|
a = array_ops.identity(1.0)
|
|
self.evaluate(a)
|
|
|
|
with self.assertRaisesRegex(errors.InvalidArgumentError,
|
|
'Could not satisfy'):
|
|
with ops.device('/device:GPU:' + str(len(logical_gpus))):
|
|
a = array_ops.identity(1.0)
|
|
self.evaluate(a)
|
|
|
|
# Modifying the GPU configuration is not supported
|
|
with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
|
|
config.set_logical_device_configuration(gpus[-1], [
|
|
context.LogicalDeviceConfiguration(memory_limit=20),
|
|
context.LogicalDeviceConfiguration(memory_limit=20)
|
|
])
|
|
|
|
with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
|
|
config.set_logical_device_configuration(gpus[-1], [
|
|
context.LogicalDeviceConfiguration(memory_limit=10),
|
|
context.LogicalDeviceConfiguration(memory_limit=10),
|
|
context.LogicalDeviceConfiguration(memory_limit=10)
|
|
])
|
|
|
|
# Setting the same GPU configuration is fine
|
|
config.set_logical_device_configuration(gpus[-1], [
|
|
context.LogicalDeviceConfiguration(memory_limit=10),
|
|
context.LogicalDeviceConfiguration(memory_limit=10)
|
|
])
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testGpuGrowth(self):
|
|
gpus = config.list_physical_devices('GPU')
|
|
self.assertNotEqual(len(gpus), 0)
|
|
|
|
self.assertIsNone(config.get_memory_growth(gpus[-1]))
|
|
for gpu in gpus:
|
|
config.set_memory_growth(gpu, True)
|
|
|
|
c = context.context().config
|
|
self.assertTrue(c.gpu_options.allow_growth)
|
|
|
|
logical_gpus = config.list_logical_devices('GPU')
|
|
self.assertTrue(len(logical_gpus), len(gpus))
|
|
|
|
# Modifying the GPU configuration is not supported
|
|
with self.assertRaisesRegex(RuntimeError, 'cannot be modified'):
|
|
for gpu in gpus:
|
|
config.set_memory_growth(gpu, False)
|
|
|
|
# Setting the same GPU configuration is fine
|
|
for gpu in gpus:
|
|
config.set_memory_growth(gpu, True)
|
|
|
|
@test_util.run_gpu_or_tpu
|
|
@reset_eager
|
|
def testGetMemoryInfoBasic(self, device_type):
|
|
with ops.device(f'{device_type}:0'):
|
|
device = array_ops.zeros([]).backing_device
|
|
info = config.get_memory_info(device)
|
|
self.assertGreater(info['current'], 0)
|
|
self.assertGreater(info['peak'], 0)
|
|
self.assertEqual(info.keys(), {'current', 'peak'})
|
|
self.assertEqual(config.get_memory_usage(device), info['current'])
|
|
|
|
@test_util.run_gpu_or_tpu
|
|
@reset_eager
|
|
def testGetMemoryUsageSubstring(self, device_type):
|
|
info = config.get_memory_info(f'{device_type}:0')
|
|
self.assertGreater(info['current'], 0)
|
|
|
|
@reset_eager
|
|
def testGetMemoryInfoCPU(self):
|
|
if test_util.IsMklEnabled():
|
|
# TODO(gzmkl) work with Google team to address design issue in allocator.h
|
|
self.skipTest('MklCPUAllocator does not throw exception. So skip test.')
|
|
|
|
with self.assertRaisesRegex(ValueError, 'Allocator stats not available'):
|
|
config.get_memory_info('CPU:0')
|
|
with self.assertRaisesRegex(ValueError, 'Allocator stats not available'):
|
|
config.get_memory_usage('CPU:0')
|
|
|
|
@reset_eager
|
|
def testGetMemoryInfoUnknownDevice(self):
|
|
with self.assertRaisesRegex(ValueError, 'No matching devices found'):
|
|
config.get_memory_info('unknown_device:0')
|
|
with self.assertRaisesRegex(ValueError, 'No matching devices found'):
|
|
config.get_memory_usage('unknown_device:0')
|
|
|
|
@reset_eager
|
|
def testGetMemoryInfoInvalidDeviceString(self):
|
|
with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
|
|
context.context().get_memory_info('GPU')
|
|
with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
|
|
context.context().get_memory_info('GPU:')
|
|
with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
|
|
context.context().get_memory_info('GPU:CPU')
|
|
|
|
@test_util.run_gpu_or_tpu
|
|
@reset_eager
|
|
def testPeakMemoryUsage(self, device_type):
|
|
device = f'{device_type}:0'
|
|
with ops.device(device):
|
|
x1 = array_ops.zeros((1000, 1000))
|
|
peak1 = config.get_memory_info(device)['peak']
|
|
self.assertGreaterEqual(peak1, 4 * 1000 * 1000)
|
|
with ops.device(device):
|
|
x2 = array_ops.ones((1000, 1000))
|
|
peak2 = config.get_memory_info(device)['peak']
|
|
self.assertGreaterEqual(peak2, peak1 + 4 * 1000 * 1000)
|
|
del x1, x2 # With CPython, causes tensor memory to be immediately freed
|
|
peak3 = config.get_memory_info(device)['peak']
|
|
self.assertGreaterEqual(peak3, peak2)
|
|
self.assertGreaterEqual(peak3, config.get_memory_info(device)['current'])
|
|
|
|
@test_util.run_gpu_or_tpu
|
|
@reset_eager
|
|
def testResetMemoryStats(self, device_type):
|
|
device = f'{device_type}:0'
|
|
with ops.device(device):
|
|
x = array_ops.zeros((1000, 1000), dtype=dtypes.float32)
|
|
config.reset_memory_stats(device)
|
|
info1 = config.get_memory_info(device)
|
|
self.assertGreaterEqual(info1['peak'], 4 * 1000 * 1000)
|
|
self.assertGreaterEqual(info1['peak'], info1['current'])
|
|
self.assertGreater(info1['current'], 0)
|
|
|
|
del x # With CPython, causes tensor memory to be immediately freed
|
|
config.reset_memory_stats(device)
|
|
info2 = config.get_memory_info(device)
|
|
self.assertLess(info2['peak'], info1['peak'])
|
|
|
|
@reset_eager
|
|
def testResetMemoryStatsCPU(self):
|
|
if test_util.IsMklEnabled():
|
|
# TODO(gzmkl) work with Google team to address design issue in allocator.h
|
|
self.skipTest('MklCPUAllocator does not throw exception. So skip test.')
|
|
|
|
with self.assertRaisesRegex(ValueError, 'Cannot reset memory stats'):
|
|
config.reset_memory_stats('CPU:0')
|
|
|
|
@reset_eager
|
|
def testResetMemoryStatsUnknownDevice(self):
|
|
with self.assertRaisesRegex(ValueError, 'No matching devices found'):
|
|
config.reset_memory_stats('unknown_device:0')
|
|
|
|
@reset_eager
|
|
def testResetMemoryStatsInvalidDeviceString(self):
|
|
with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
|
|
context.context().reset_memory_stats('GPU')
|
|
with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
|
|
context.context().reset_memory_stats('GPU:')
|
|
with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
|
|
context.context().reset_memory_stats('GPU:CPU')
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testGpuInvalidConfig(self):
|
|
gpus = config.list_physical_devices('GPU')
|
|
self.assertNotEqual(len(gpus), 0)
|
|
|
|
if len(gpus) > 1:
|
|
# Assert if other GPUs were not configured
|
|
config.set_memory_growth(gpus[0], True)
|
|
with self.assertRaisesRegex(ValueError, 'cannot differ'):
|
|
c = context.context().config
|
|
|
|
# If we limit visibility to GPU 0, growth is fine
|
|
config.set_visible_devices(gpus[0], 'GPU')
|
|
c = context.context().config
|
|
self.assertTrue(c.gpu_options.allow_growth)
|
|
|
|
# Default setting for second GPU is False and works if we set visibility
|
|
config.set_visible_devices(gpus[1], 'GPU')
|
|
c = context.context().config
|
|
self.assertFalse(c.gpu_options.allow_growth)
|
|
|
|
# Growth now fails because all the GPUs are visible and not the same
|
|
config.set_visible_devices(gpus, 'GPU')
|
|
with self.assertRaisesRegex(ValueError, 'cannot differ'):
|
|
c = context.context().config
|
|
|
|
for gpu in gpus:
|
|
config.set_memory_growth(gpu, True)
|
|
|
|
c = context.context().config
|
|
self.assertTrue(c.gpu_options.allow_growth)
|
|
|
|
with self.assertRaisesRegex(ValueError, 'memory limit'):
|
|
config.set_logical_device_configuration(gpus[-1], [
|
|
context.LogicalDeviceConfiguration(),
|
|
context.LogicalDeviceConfiguration()
|
|
])
|
|
|
|
self.assertIsNone(config.get_logical_device_configuration(gpus[-1]))
|
|
config.set_logical_device_configuration(gpus[-1], [
|
|
context.LogicalDeviceConfiguration(memory_limit=10),
|
|
context.LogicalDeviceConfiguration(memory_limit=10)
|
|
])
|
|
|
|
c = context.context().config
|
|
self.assertFalse(c.gpu_options.allow_growth)
|
|
|
|
with self.assertRaisesRegex(ValueError, 'virtual devices'):
|
|
config.set_memory_growth(gpus[-1], False)
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testRemote(self):
|
|
gpus = config.list_logical_devices('GPU')
|
|
self.assertNotEqual(len(gpus), 0)
|
|
|
|
context.ensure_initialized()
|
|
|
|
gpus = config.list_logical_devices('GPU')
|
|
self.assertNotEqual(len(gpus), 0)
|
|
for gpu in gpus:
|
|
self.assertIsNotNone(gpu.name)
|
|
|
|
context.ensure_initialized()
|
|
|
|
job_name = 'test'
|
|
cluster_def = cluster_pb2.ClusterDef()
|
|
job_def = cluster_def.job.add()
|
|
job_def.name = job_name
|
|
job_def.tasks[0] = 'localhost:0'
|
|
|
|
server_def = tensorflow_server_pb2.ServerDef(
|
|
cluster=cluster_def, job_name=job_name, task_index=0, protocol='grpc')
|
|
|
|
context.set_server_def(server_def)
|
|
|
|
gpus = config.list_logical_devices('GPU')
|
|
for gpu in gpus:
|
|
self.assertIsNotNone(gpu.name)
|
|
|
|
@reset_eager
|
|
def testV1CompatibilityDummyInvisibleDeviceList(self):
|
|
gpus = config.list_physical_devices('GPU')
|
|
if gpus:
|
|
self.skipTest('Test requires no GPUs')
|
|
|
|
# Ensure GPU options left untouched on CPU only environments
|
|
context.context()._physical_devices = None
|
|
context.context()._config = config_pb2.ConfigProto(
|
|
gpu_options=config_pb2.GPUOptions(visible_device_list='0'))
|
|
new_config = context.context().config
|
|
self.assertEqual(new_config.gpu_options.visible_device_list, '0')
|
|
|
|
@test_util.run_gpu_only
|
|
@reset_eager
|
|
def testV1Compatibility(self):
|
|
# Ensure we set 1 CPU by default
|
|
context.context()._config = config_pb2.ConfigProto()
|
|
new_config = context.context().config
|
|
self.assertEqual(new_config.device_count['CPU'], 1)
|
|
context.context()._physical_devices = None
|
|
|
|
# Ensure CPU is split
|
|
context.context()._config = config_pb2.ConfigProto(device_count={'CPU': 2})
|
|
new_config = context.context().config
|
|
self.assertEqual(new_config.device_count['CPU'], 2)
|
|
context.context()._physical_devices = None
|
|
|
|
# Handle empty visible device list
|
|
context.context()._config = config_pb2.ConfigProto(
|
|
gpu_options=config_pb2.GPUOptions(visible_device_list=''))
|
|
gpus = config.list_physical_devices('GPU')
|
|
gpu_count = len(gpus)
|
|
new_config = context.context().config
|
|
self.assertEqual(new_config.gpu_options.visible_device_list,
|
|
','.join(str(i) for i in range(len(gpus))))
|
|
context.context()._physical_devices = None
|
|
|
|
# Handle invalid visible device list
|
|
context.context()._config = config_pb2.ConfigProto(
|
|
gpu_options=config_pb2.GPUOptions(visible_device_list=str(gpu_count)))
|
|
with self.assertRaisesRegex(ValueError, 'Invalid visible device index'):
|
|
gpus = config.list_physical_devices('GPU')
|
|
new_config = context.context().config
|
|
context.context()._physical_devices = None
|
|
|
|
# Handle single visible device list
|
|
context.context()._config = config_pb2.ConfigProto(
|
|
gpu_options=config_pb2.GPUOptions(
|
|
visible_device_list=str(gpu_count - 1)))
|
|
gpus = config.list_physical_devices('GPU')
|
|
new_config = context.context().config
|
|
self.assertEqual(new_config.gpu_options.visible_device_list,
|
|
str(gpu_count - 1))
|
|
context.context()._physical_devices = None
|
|
|
|
def testConfigureCollectiveOps(self):
|
|
context.context().configure_collective_ops(
|
|
collective_leader='/job:worker/replica:0/task:0',
|
|
scoped_allocator_enabled_ops=('CollectiveReduce',),
|
|
use_nccl_communication=False,
|
|
device_filters=['/job:worker/task:1'])
|
|
new_config = context.context().config
|
|
|
|
# Verify group leader
|
|
self.assertEqual('/job:worker/replica:0/task:0',
|
|
new_config.experimental.collective_group_leader)
|
|
|
|
# Verify device filters.
|
|
self.assertEqual(['/job:worker/task:1'], new_config.device_filters)
|
|
|
|
# Verify rewrite options.
|
|
new_rewrite_options = new_config.graph_options.rewrite_options
|
|
self.assertEqual(rewriter_config_pb2.RewriterConfig.ON,
|
|
new_rewrite_options.scoped_allocator_optimization)
|
|
self.assertEqual(['CollectiveReduce'],
|
|
new_rewrite_options.scoped_allocator_opts.enable_op)
|
|
|
|
def testDeterminism(self):
|
|
# This does not test any ops are deterministic, because that is tested by
|
|
# many kernel tests.
|
|
try:
|
|
config.disable_op_determinism()
|
|
self.assertFalse(config.is_op_determinism_enabled())
|
|
config.enable_op_determinism()
|
|
self.assertTrue(config.is_op_determinism_enabled())
|
|
finally:
|
|
config.disable_op_determinism()
|
|
|
|
|
|
class TensorFloat32Test(test.TestCase):
|
|
|
|
def tearDown(self):
|
|
super(TensorFloat32Test, self).tearDown()
|
|
config.enable_tensor_float_32_execution(True)
|
|
|
|
def test_tensor_float_32_global_variable(self):
|
|
self.assertTrue(config.tensor_float_32_execution_enabled())
|
|
self.assertTrue(test_ops.is_tensor_float32_enabled())
|
|
config.enable_tensor_float_32_execution(False)
|
|
self.assertFalse(config.tensor_float_32_execution_enabled())
|
|
self.assertFalse(test_ops.is_tensor_float32_enabled())
|
|
config.enable_tensor_float_32_execution(True)
|
|
self.assertTrue(config.tensor_float_32_execution_enabled())
|
|
self.assertTrue(test_ops.is_tensor_float32_enabled())
|
|
|
|
def _skip_if_tensor_float_32_unsupported(self):
|
|
if not test_util.is_gpu_available(
|
|
cuda_only=True, min_cuda_compute_capability=(8, 0)):
|
|
self.skipTest('TensorFloat-32 requires an NVIDIA GPU with compute '
|
|
'capability of at least 8.0')
|
|
|
|
# Size of each dimension of matrices to test. cuBLAS does not use TF32 for
|
|
# small matrices, so we must choose a large enough size to cause TF32 to be
|
|
# used.
|
|
DIM = 2 ** 10
|
|
|
|
def test_tensor_float_32_enabled(self):
|
|
self._skip_if_tensor_float_32_unsupported()
|
|
self.assertTrue(config.tensor_float_32_execution_enabled())
|
|
|
|
x = array_ops.fill((self.DIM, self.DIM), 1 + 2**-12)
|
|
y = array_ops.ones((self.DIM, self.DIM))
|
|
out = math_ops.matmul(x, y)
|
|
# In TensorFloat-32, each element of x is rounded to 1, so each output
|
|
# element should be self.DIM.
|
|
expected = array_ops.fill((self.DIM, self.DIM), float(self.DIM))
|
|
self.assertAllEqual(out, expected)
|
|
|
|
def test_tensor_float_32_disabled(self):
|
|
self._skip_if_tensor_float_32_unsupported()
|
|
self.assertTrue(config.tensor_float_32_execution_enabled())
|
|
config.enable_tensor_float_32_execution(False)
|
|
self.assertFalse(config.tensor_float_32_execution_enabled())
|
|
|
|
x = array_ops.fill((self.DIM, self.DIM), 1 + 2**-12)
|
|
y = array_ops.ones((self.DIM, self.DIM))
|
|
out = math_ops.matmul(x, y)
|
|
expected = array_ops.fill((self.DIM, self.DIM), self.DIM * (1 + 2**-12))
|
|
self.assertAllClose(out, expected, rtol=2**-13, atol=0)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
ops.enable_eager_execution()
|
|
test.main()
|