Skip to content

Commit

Permalink
Use a cache variable to store local devices.
Browse files Browse the repository at this point in the history
  • Loading branch information
datumbox committed Nov 4, 2017
1 parent 9605ed7 commit 8201921
Showing 1 changed file with 9 additions and 1 deletion.
10 changes: 9 additions & 1 deletion keras/backend/tensorflow_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,11 @@
# Change its value via `manual_variable_initialization(value)`.
_MANUAL_VAR_INIT = False

# This list holds the available devices.
# It is populated when `_get_available_gpus()` is called for the first time.
# We assume our devices don't change during our lifetime.
_LOCAL_DEVICES = None


def get_uid(prefix=''):
"""Get the uid for the default graph.
Expand Down Expand Up @@ -246,7 +251,10 @@ def _get_available_gpus():
# Returns
A list of available GPU devices.
"""
return [x.name for x in get_session().list_devices() if x.device_type == 'GPU']
global _LOCAL_DEVICES
if _LOCAL_DEVICES is None:
_LOCAL_DEVICES = get_session().list_devices()
return [x.name for x in _LOCAL_DEVICES if x.device_type == 'GPU']


def _has_nchw_support():
Expand Down

0 comments on commit 8201921

Please sign in to comment.