fix(encoder): fix gpu limitation in inception

gnes-ai · Jul 25, 2019 · e644e39 · e644e39
1 parent 89d8b70
commit e644e39
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/gnes/encoder/image/inception.py b/gnes/encoder/image/inception.py
@@ -19,7 +19,7 @@
 from PIL import Image
 
 from ..base import BaseImageEncoder
-from ...helper import batching, batch_iterator
+from ...helper import batching, batch_iterator, get_first_available_gpu
 
 
 class TFInceptionEncoder(BaseImageEncoder):
@@ -42,7 +42,8 @@ def post_init(self):
         import tensorflow as tf
         from .inception_cores.inception_v4 import inception_v4
         from .inception_cores.inception_utils import inception_arg_scope
-
+        import os
+        os.environ['CUDA_VISIBLE_DEVICES'] = get_first_available_gpu()
         g = tf.Graph()
         with g.as_default():
             arg_scope = inception_arg_scope()

diff --git a/gnes/helper.py b/gnes/helper.py
@@ -55,9 +55,9 @@ def get_first_available_gpu():
             return r[0]
         raise ValueError
     except ImportError:
-        return 0
+        return -1
     except ValueError:
-        return 0
+        return -1
 
 
 class FileLock: