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: 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 _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 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)