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