Merge pull request thmoa#6 from rcrvano/master

set tf.config allow_growth
grll · Apr 29, 2019 · 4bab53b · 4bab53b
2 parents 2b23cab + 27ef61e
commit 4bab53b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/infer_batch.py b/infer_batch.py
@@ -1,6 +1,8 @@
 import os
 import csv
 import argparse
+import tensorflow as tf
+import keras.backend as K
 
 from glob import glob
 
@@ -9,6 +11,8 @@
 
 
 def main(weights, num, batch_file, opt_pose_steps, opt_shape_steps):
+    K.set_session(tf.Session(config=tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))))
+
     model = Octopus(num=num)
 
     with open(batch_file, 'rb') as f:

diff --git a/infer_single.py b/infer_single.py
@@ -1,5 +1,7 @@
 import os
 import argparse
+import tensorflow as tf
+import keras.backend as K
 
 from glob import glob
 
@@ -14,6 +16,8 @@ def main(weights, name, segm_dir, pose_dir, out_dir, opt_pose_steps, opt_shape_s
     if len(segm_files) != len(pose_files) or len(segm_files) == len(pose_files) == 0:
         exit('Inconsistent input.')
 
+    K.set_session(tf.Session(config=tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True))))
+
     model = Octopus(num=len(segm_files))
     model.load(weights)