def setUp(self): super(EagerTensorCacheTest, self).setUp() context._reset_context() configure_virtual_cpus()
def setUp(self): super(TFETest, self).setUp() context._reset_context() configure_virtual_cpus()
def reset_context():
context._reset_context() # pylint: disable=protected-access
def setUp(self):
super(DeviceUtilTest, self).setUp()
context._reset_context() # pylint: disable=protected-access
def tearDown(self):
super(RemoteAsyncTest, self).tearDown()
# Reset the context to avoid polluting other test cases.
context._reset_context()
def tearDownClass(cls): super(RemoteExecutionWithoutLazyRemoteInputsCopyTest, cls).tearDownClass() context._reset_context() context.context().lazy_remote_inputs_copy = True
def setUpClass(cls): super(DynamicClusterWithoutLazyRemoteInputsCopyTest, cls).setUpClass() context._reset_context() context.context().lazy_remote_inputs_copy = False
def test_save_variable_devices(self, save_devices, meta_graph_only):
context._reset_context()
cpus = context.context().list_physical_devices("CPU")
if len(cpus) == 1:
context.context().set_logical_device_configuration(
cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
context.ensure_initialized()
root = tracking.AutoTrackable()
with ops.device("CPU:0"):
root.v0 = variables.Variable(1., name="v0")
with ops.device("CPU:1"):
root.v1 = variables.Variable(1., name="v1")
options = save_options.SaveOptions(
experimental_variable_policy=save_devices)
file_name = os.path.join(self.get_temp_dir(), "saved_model")
if meta_graph_only:
save.export_meta_graph(obj=root,
filename=file_name,
options=options)
else:
save.save(obj=root, export_dir=file_name, options=options)
meta = None
if meta_graph_only:
meta = meta_graph.read_meta_graph_file(file_name)
else:
meta = loader_impl.parse_saved_model(file_name).meta_graphs[0]
# Check devices in meta graph nodes.
graph_def = meta.graph_def
v0 = next((n for n in graph_def.node if n.name == "v0"), None)
v1 = next((n for n in graph_def.node if n.name == "v1"), None)
self.assertIsNotNone(v0)
self.assertIsNotNone(v1)
if save_devices == save_options.VariablePolicy.SAVE_VARIABLE_DEVICES:
self.assertIn("CPU:0", v0.device)
self.assertIn("CPU:1", v1.device)
else:
self.assertEmpty(v0.device)
self.assertEmpty(v1.device)
# Check devices in object graph nodes.
object_graph_def = meta.object_graph_def
v0 = next((n.variable for n in object_graph_def.nodes
if n.HasField("variable") and n.variable.name == "v0"),
None)
v1 = next((n.variable for n in object_graph_def.nodes
if n.HasField("variable") and n.variable.name == "v1"),
None)
self.assertIsNotNone(v0)
self.assertIsNotNone(v1)
if save_devices == save_options.VariablePolicy.SAVE_VARIABLE_DEVICES:
self.assertIn("CPU:0", v0.device)
self.assertIn("CPU:1", v1.device)
else:
self.assertEmpty(v0.device)
self.assertEmpty(v1.device)
def setUp(self):
context._reset_context() # pylint: disable=protected-access
super(CollectiveOpTest, self).setUp()
def testLogicalCPUs(self):
context._reset_context()
test_util.set_logical_devices_to_at_least('CPU', 3)
cpu_device = config.list_physical_devices('CPU')[0]
self.assertLen(config.get_logical_device_configuration(cpu_device), 3)
def setUp(self):
super(SoftDevicePlacementTest, self).setUp()
context._reset_context()
config.set_soft_device_placement(enabled=True)
context.context().log_device_placement = True
def testDLPackFromWithoutContextInitialization(self):
tf_tensor = constant_op.constant(1)
dlcapsule = dlpack.to_dlpack(tf_tensor)
# Resetting the context doesn't cause an error.
context._reset_context()
_ = dlpack.from_dlpack(dlcapsule)
def setUp(self):
context._reset_context() # pylint: disable=protected-access
# Need to call set_virtual_cpus_to_at_least() in setUp with the maximum
# value needed in any test.
strategy_combinations.set_virtual_cpus_to_at_least(3)
super(VirtualDevicesTest, self).setUp()
def tearDown(self):
super().tearDown()
# reset context to disconnect from the cluster.
context._reset_context()
def tearDownClass(cls): super(DynamicClusterWithoutLazyRemoteInputsCopyTest, cls).tearDownClass() context._reset_context() context.context().lazy_remote_inputs_copy = True
def tearDown(self):
super().tearDown()
context._reset_context()
def _setup_context():
context._reset_context()
test_util.set_logical_devices_to_at_least('CPU', 4)
context.ensure_initialized()
def _configure_tpu_runtime():
was_enabled = context.is_tfrt_enabled()
if ("tpu_use_tfrt" in flags.FLAGS and flags.FLAGS["tpu_use_tfrt"].value):
tfrt_utils.set_tfrt_enabled(True)
if not was_enabled:
context._reset_context() # pylint:disable=protected-access
def setUp(self):
super(TFETest, self).setUp()
ops.device(None).__enter__()
context._reset_context()
configure_virtual_cpus()
def tearDown(self): super(DynamicClusterTest, self).tearDown() ops.device(None).__enter__() context._reset_context()
def setUpClass(cls): super(RemoteExecutionWithoutLazyRemoteInputsCopyTest, cls).setUpClass() context._reset_context() context.context().lazy_remote_inputs_copy = False