From 64acbc93ad67a1dd2869fa4b5cb5ae8a04ac3b34 Mon Sep 17 00:00:00 2001
From: Tiago Freitas Pereira <tiagofrepereira@gmail.com>
Date: Sun, 23 Oct 2016 18:14:27 +0200
Subject: [PATCH] Set environment variable to map the devices
---
bob/learn/tensorflow/trainers/Trainer.py | 5 ++++-
1 file changed, 4 insertions(+), 1 deletion(-)
diff --git a/bob/learn/tensorflow/trainers/Trainer.py b/bob/learn/tensorflow/trainers/Trainer.py
index d3500ab6..ecd30aa1 100644
--- a/bob/learn/tensorflow/trainers/Trainer.py
+++ b/bob/learn/tensorflow/trainers/Trainer.py
@@ -14,6 +14,7 @@ from tensorflow.core.framework import summary_pb2
import time
from bob.learn.tensorflow.datashuffler.OnlineSampling import OnLineSampling
+os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,3,0"
logger = bob.core.log.setup("bob.learn.tensorflow")
@@ -276,7 +277,9 @@ class Trainer(object):
# Training
hdf5 = bob.io.base.HDF5File(os.path.join(self.temp_dir, 'model.hdf5'), 'w')
- with tf.Session() as session:
+ config = tf.ConfigProto(log_device_placement=True)
+ config.gpu_options.allow_growth = True
+ with tf.Session(config=True) as session:
tf.initialize_all_variables().run()
--
GitLab