def _initialize_devices(self):
"""Helper to initialize devices."""
# Store list of devices
self._context_devices = []
device_list = pywrap_tensorflow.TFE_ContextListDevices(
self._context_handle)
try:
self._num_gpus = 0
for i in range(pywrap_tensorflow.TF_DeviceListCount(device_list)):
dev_name = pywrap_tensorflow.TF_DeviceListName(device_list, i)
self._context_devices.append(pydev.canonical_name(dev_name))
dev_type = pywrap_tensorflow.TF_DeviceListType(device_list, i)
if dev_type == "GPU":
self._num_gpus += 1
finally:
pywrap_tensorflow.TF_DeleteDeviceList(device_list)
def testInvalidDeviceNumber(self):
opts = tf_session.TF_NewSessionOptions()
with errors.raise_exception_on_not_ok_status() as status:
c_session = tf_session.TF_NewSession(
ops.get_default_graph()._c_graph, opts, status)
raw_device_list = tf_session.TF_SessionListDevices(
c_session, status)
size = tf_session.TF_DeviceListCount(raw_device_list)
# Test that invalid device numbers return -1 rather than a Swig-wrapped
# pointer.
status_no_exception = c_api_util.ScopedTFStatus()
memory = tf_session.TF_DeviceListMemoryBytes(
raw_device_list, size, status_no_exception)
self.assertEqual(memory, -1)
tf_session.TF_DeleteDeviceList(raw_device_list)
with errors.raise_exception_on_not_ok_status() as status:
tf_session.TF_CloseSession(c_session, status)
def list_devices(session):
with errors.raise_exception_on_not_ok_status() as status:
if session._created_with_new_api:
raw_device_list = tf_session.TF_SessionListDevices(
session._session, status)
else:
raw_device_list = tf_session.TF_DeprecatedSessionListDevices(
session._session, status)
device_list = []
size = tf_session.TF_DeviceListCount(raw_device_list)
for i in range(size):
name = tf_session.TF_DeviceListName(raw_device_list, i, status)
device_type = tf_session.TF_DeviceListType(raw_device_list, i, status)
memory = 0 # tf_session.TF_DeviceListMemoryBytes(raw_device_list, i, status)
device_list.append(_DeviceAttributes(name, device_type, memory))
tf_session.TF_DeleteDeviceList(raw_device_list)
return device_list
def _initialize_handle_and_devices(self):
"""Initialize handle and devices."""
with self._initialize_lock:
if self._context_handle is not None:
return
assert self._context_devices is None
opts = pywrap_tensorflow.TFE_NewContextOptions()
try:
with errors.raise_exception_on_not_ok_status() as status:
if self._config is not None:
config_str = self._config.SerializeToString()
pywrap_tensorflow.TFE_ContextOptionsSetConfig(
opts, config_str, len(config_str), status)
if self._device_policy is not None:
pywrap_tensorflow.TFE_ContextOptionsSetDevicePlacementPolicy(
opts, self._device_policy)
if self._execution_mode == ASYNC:
pywrap_tensorflow.TFE_ContextOptionsSetAsync(True)
self._context_handle = pywrap_tensorflow.TFE_NewContext(
opts, status)
finally:
pywrap_tensorflow.TFE_DeleteContextOptions(opts)
# Store list of devices
self._context_devices = []
with errors.raise_exception_on_not_ok_status() as status:
device_list = pywrap_tensorflow.TFE_ContextListDevices(
self._context_handle, status)
try:
self._num_gpus = 0
for i in range(
pywrap_tensorflow.TF_DeviceListCount(device_list)):
with errors.raise_exception_on_not_ok_status() as status:
dev_name = pywrap_tensorflow.TF_DeviceListName(
device_list, i, status)
self._context_devices.append(
pydev.canonical_name(dev_name))
with errors.raise_exception_on_not_ok_status() as status:
dev_type = pywrap_tensorflow.TF_DeviceListType(
device_list, i, status)
if dev_type == "GPU":
self._num_gpus += 1
finally:
pywrap_tensorflow.TF_DeleteDeviceList(device_list)
def __init__(self):
self._eager_context = _EagerContext()
# Create a handle
opts = pywrap_tensorflow.TF_NewSessionOptions(
target=compat.as_bytes(""), config=None)
with errors.raise_exception_on_not_ok_status() as status:
self._handle = pywrap_tensorflow.TFE_NewContext(opts, status)
pywrap_tensorflow.TF_DeleteSessionOptions(opts)
# Store list of devices
self._devices = []
with errors.raise_exception_on_not_ok_status() as status:
device_list = pywrap_tensorflow.TFE_ContextListDevices(
self._handle, status)
try:
for i in range(pywrap_tensorflow.TF_DeviceListCount(device_list)):
with errors.raise_exception_on_not_ok_status() as status:
dev_name = pywrap_tensorflow.TF_DeviceListName(
device_list, i, status)
self._devices.append(pydev.canonical_name(dev_name))
finally:
pywrap_tensorflow.TF_DeleteDeviceList(device_list)
self._summary_writer_resource = None
def _initialize_handle_and_devices(self):
"""Initialize handle and devices."""
with self._initialize_lock:
if self._context_handle is not None:
return
assert self._context_devices is None
opts = pywrap_tensorflow.TF_NewSessionOptions(
target=compat.as_bytes(""), config=self._config)
with errors.raise_exception_on_not_ok_status() as status:
self._context_handle = pywrap_tensorflow.TFE_NewContext(opts, status)
pywrap_tensorflow.TF_DeleteSessionOptions(opts)
# Store list of devices
self._context_devices = []
with errors.raise_exception_on_not_ok_status() as status:
device_list = pywrap_tensorflow.TFE_ContextListDevices(
self._context_handle, status)
try:
for i in range(pywrap_tensorflow.TF_DeviceListCount(device_list)):
with errors.raise_exception_on_not_ok_status() as status:
dev_name = pywrap_tensorflow.TF_DeviceListName(
device_list, i, status)
self._context_devices.append(pydev.canonical_name(dev_name))
finally:
pywrap_tensorflow.TF_DeleteDeviceList(device_list)
