Skip to content
Permalink
Browse files

fix(encoder): fix gpu limitation in inception

  • Loading branch information...
Larryjianfeng committed Jul 25, 2019
1 parent 89d8b70 commit e644e3916222d18a9ebe0329fd600b6f776fb69a
Showing with 5 additions and 4 deletions.
  1. +3 −2 gnes/encoder/image/inception.py
  2. +2 −2 gnes/helper.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()
@@ -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:

0 comments on commit e644e39

Please sign in to comment.
You can’t perform that action at this time.