def setUp(self):
super(_VirtualDeviceTestCase, self).setUp()
ctx = context.context()
if ctx.list_physical_devices("TPU"):
self.device_type = "TPU"
tpu_strategy_util.initialize_tpu_system()
elif ctx.list_physical_devices("GPU"):
self.device_type = "GPU"
gpus = ctx.list_physical_devices(self.device_type)
ctx.set_logical_device_configuration(gpus[0], [
context.LogicalDeviceConfiguration(memory_limit=100),
context.LogicalDeviceConfiguration(memory_limit=100),
])
else:
self.device_type = "CPU"
cpus = ctx.list_physical_devices("CPU")
ctx.set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration(),
])
self.device = parallel_device.ParallelDevice(components=[
"/job:localhost/device:{}:0".format(self.device_type),
self.device_type + ":1"
])
self.assertIn(self.device_type + ":0", self.device.components[0])
self.assertIn(self.device_type + ":1", self.device.components[1])
def test_one_replica_eager_control_flow(self):
device = parallel_device.ParallelDevice(components=[
"/job:localhost/device:{}:0".format(self.device_type),
])
x = constant_op.constant([2, 3, 4])
with device:
x = device.pack([x])
if math_ops.reduce_any(math_ops.equal(x, constant_op.constant(4))):
y = constant_op.constant(1)
else:
y = constant_op.constant(2)
self.assertAllEqual([1], device.unpack(y))
def setUp(self):
super(_VirtualDeviceTestCase, self).setUp()
cpus = context.context().list_physical_devices("CPU")
# Set 4 virtual CPUs
context.context().set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
self.device = parallel_device.ParallelDevice(
components=["/job:localhost/device:CPU:0", "CPU:1"])
self.assertIn("CPU:0", self.device.components[0])
self.assertIn("CPU:1", self.device.components[1])
def setUp(self):
super(_VirtualDeviceTestCase, self).setUp()
cpus = context.context().list_physical_devices("CPU")
# Set 4 virtual CPUs
context.context().set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
# TODO(allenl): Make CPU:0 and CPU:1 work (right now "CPU:1" soft-places
# onto CPU:0, which seems wrong).
components = [
"/job:localhost/replica:0/task:0/device:CPU:0",
"/job:localhost/replica:0/task:0/device:CPU:1"
]
self.device = parallel_device.ParallelDevice(components)
