def device_function(self, op):
"""Chose a device for `op`.
Args:
op: an `Operation`.
Returns:
The device to use for the `Operation`.
"""
if not self._merge_devices and op.device:
return op.device
current_device = pydev.from_string(op.device or "")
spec = pydev.Device()
if self._ps_tasks and self._ps_device:
node_def = op if isinstance(op, graph_pb2.NodeDef) else op.node_def
if node_def.op in self._ps_ops:
device_string = "%s/task:%d" % (self._ps_device, self._next_task())
if self._merge_devices:
spec = pydev.from_string(device_string)
spec.merge_from(current_device)
return spec.to_string()
else:
return device_string
if self._worker_device:
if not self._merge_devices:
return self._worker_device
spec = pydev.from_string(self._worker_device)
if not self._merge_devices:
return ""
spec.merge_from(current_device)
return spec.to_string()
def device_function(self, op):
"""Chose a device for `op`.
Args:
op: an `Operation`.
Returns:
The device to use for the `Operation`.
"""
if not self._merge_devices and op.device:
return op.device
current_device = pydev.from_string(op.device or "")
spec = pydev.Device()
if self._ps_tasks and self._ps_device:
node_def = op if isinstance(op, graph_pb2.NodeDef) else op.node_def
if node_def.op in self._ps_ops:
device_string = "%s/task:%d" % (self._ps_device,
self._next_task())
if self._merge_devices:
spec = pydev.from_string(device_string)
spec.merge_from(current_device)
return spec.to_string()
else:
return device_string
if self._worker_device:
if not self._merge_devices:
return self._worker_device
spec = pydev.from_string(self._worker_device)
if not self._merge_devices:
return ""
spec.merge_from(current_device)
return spec.to_string()
def testFromString(self):
d = device.from_string("/job:foo/replica:0")
self.assertEquals("/job:foo/replica:0", d.to_string())
with self.assertRaises(Exception) as e:
d = device.from_string("/job:muu/gpu:2/cpu:0")
self.assertTrue("Cannot specify multiple device" in e.exception.message)
d = device.from_string("/job:foo/replica:0/task:3/cpu:*")
self.assertEquals(None, d.device_index)
d = device.from_string("/job:foo/replica:0/task:3/gpu:7")
self.assertEquals(7, d.device_index)
d = device.from_string("/job:foo/replica:0/task:3/device:GPU:7")
self.assertEquals(7, d.device_index)
def testFromString(self):
d = device.from_string("/job:foo/replica:0")
self.assertEquals("/job:foo/replica:0", d.to_string())
with self.assertRaises(Exception) as e:
d = device.from_string("/job:muu/gpu:2/cpu:0")
self.assertTrue("Cannot specify multiple device" in str(e.exception))
d = device.from_string("/job:foo/replica:0/task:3/cpu:*")
self.assertEquals(None, d.device_index)
d = device.from_string("/job:foo/replica:0/task:3/gpu:7")
self.assertEquals(7, d.device_index)
d = device.from_string("/job:foo/replica:0/task:3/device:GPU:7")
self.assertEquals(7, d.device_index)
def pin_variables_on_cpu(op):
"""Returns a CPU device for Variable nodes if the device is not specified.
Args:
op: The ops.Operation object describing the node for which a device
should be chosen. The op.device field is respected.
Returns:
A device containing "/device:CPU:0" if the node is related to a variable.
"""
device = op.device if op.device is not None else ""
dev = pydev.from_string(device)
# If a device type exists already, do not override.
if dev.device_type:
return device
if isinstance(op, ops.Operation):
node_def = op.node_def
else:
assert isinstance(op, graph_pb2.NodeDef)
node_def = op
if _is_variable_op(node_def.op):
return set_cpu0(device)
return device
def pin_variables_on_cpu(op):
"""Returns a CPU device for Variable nodes if the device is not specified.
Args:
op: The ops.Operation object describing the node for which a device
should be chosen. The op.device field is respected.
Returns:
A device containing "/device:CPU:0" if the node is related to a variable.
"""
device = op.device if op.device is not None else ""
dev = pydev.from_string(device)
# If a device type exists already, do not override.
if dev.device_type:
return device
if isinstance(op, ops.Operation):
node_def = op.node_def
else:
assert isinstance(op, graph_pb2.NodeDef)
node_def = op
if _is_variable_op(node_def.op):
return set_cpu0(device)
return device
def set_cpu0(device_string):
"""Creates a new device string based on `device_string' but using /CPU:0.
If the device is already on /CPU:0, this is a no-op.
Args:
device_string: A device string.
Returns:
A device string.
"""
parsed_device = pydev.from_string(device_string)
parsed_device.device_type = "CPU"
parsed_device.device_index = 0
return parsed_device.to_string()
def pin_to_cpu(op):
"""Returns a CPU device for the given node."""
device = op.device if op.device is not None else ""
dev = pydev.from_string(device)
if not dev.device_type:
return set_cpu0(device)
if dev.device_type == "CPU":
return device
logging.info("Operation %s has been assigned to a non-CPU (%s), so "
"it will not be pinned to the CPU.", op.name, dev.device_type)
return device
def pin_to_cpu(op):
"""Returns a CPU device for the given node."""
device = op.device if op.device is not None else ""
dev = pydev.from_string(device)
if not dev.device_type:
return set_cpu0(device)
if dev.device_type == "CPU":
return device
logging.info("Operation %s has been assigned to a non-CPU (%s), so "
"it will not be pinned to the CPU.", op.name, dev.device_type)
return device
def set_cpu0(device_string):
"""Creates a new device string based on `device_string' but using /CPU:0.
If the device is already on /CPU:0, this is a no-op.
Args:
device_string: A device string.
Returns:
A device string.
"""
parsed_device = pydev.from_string(device_string)
parsed_device.device_type = "CPU"
parsed_device.device_index = 0
return parsed_device.to_string()
def testMerge(self):
d = device.from_string("/job:foo/replica:0")
self.assertEquals("/job:foo/replica:0", d.to_string())
d.merge_from(device.from_string("/task:1/gpu:2"))
self.assertEquals("/job:foo/replica:0/task:1/device:GPU:2", d.to_string())
d = device.Device()
d.merge_from(device.from_string("/task:1/cpu:0"))
self.assertEquals("/task:1/device:CPU:0", d.to_string())
d.merge_from(device.from_string("/job:boo/gpu:0"))
self.assertEquals("/job:boo/task:1/device:GPU:0", d.to_string())
d.merge_from(device.from_string("/job:muu/cpu:2"))
self.assertEquals("/job:muu/task:1/device:CPU:2", d.to_string())
d.merge_from(device.from_string("/job:muu/device:MyFunnyDevice:2"))
self.assertEquals("/job:muu/task:1/device:MyFunnyDevice:2", d.to_string())
def testMerge(self):
d = device.from_string("/job:foo/replica:0")
self.assertEquals("/job:foo/replica:0", d.to_string())
d.merge_from(device.from_string("/task:1/gpu:2"))
self.assertEquals("/job:foo/replica:0/task:1/device:GPU:2",
d.to_string())
d = device.Device()
d.merge_from(device.from_string("/task:1/cpu:0"))
self.assertEquals("/task:1/device:CPU:0", d.to_string())
d.merge_from(device.from_string("/job:boo/gpu:0"))
self.assertEquals("/job:boo/task:1/device:GPU:0", d.to_string())
d.merge_from(device.from_string("/job:muu/cpu:2"))
self.assertEquals("/job:muu/task:1/device:CPU:2", d.to_string())
d.merge_from(device.from_string("/job:muu/device:MyFunnyDevice:2"))
self.assertEquals("/job:muu/task:1/device:MyFunnyDevice:2",
d.to_string())
