Run hyperparameter optimization on parallel gpus using tensorflow - multithreading

Question

I have a training function that trains a tf model end-to-end here (contrived for illustration only):
def opt_fx(params, gpu):
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
sess = tf.Session()
# Run some training on a particular gpu...
sess.run(...)
I want to run hyperparameter optimization across 20 trials using a model per gpu:
from threading import Thread
exp_trials = list(hyperparams.trials(num=20))
train_threads = []
for gpu_num, trial_params in zip(['0', '1', '2', '3']*5, exp_trials):
t = Thread(target=opt_fx, args=(trial_params, gpu_num,))
train_threads.append(t)
# Start the threads, and block on their completion.
for t in train_threads:
t.start()
for t in train_threads:
t.join()
This fails however... what's the right way to do this?

Answer 1

I'm not sure if this is the best approach, but what I ended up doing is define a graph per device and train each one in a separate session. This can be parallelized. I tried to reuse the graph in separate devices, but that didn't work. Here's how my version looks like in code (a complete example):
import threading
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
# Get the data
mnist = input_data.read_data_sets("data/mnist", one_hot=True)
train_x_all = mnist.train.images
train_y_all = mnist.train.labels
test_x = mnist.test.images
test_y = mnist.test.labels
# Define the graphs per device
devices = ['/gpu', '/cpu'] # just one GPU on this machine...
learning_rates = [0.01, 0.03]
jobs = []
for device, learning_rate in zip(devices, learning_rates):
with tf.Graph().as_default() as graph:
x = tf.placeholder(tf.float32, [None, 784], name='x')
y = tf.placeholder(tf.float32, [None, 10], name='y')
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))
pred = tf.nn.softmax(tf.matmul(x, W) + b)
accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(pred, 1), tf.argmax(y, 1)), tf.float32))
cost = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred), reduction_indices=1), name='cost')
optimize = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost, name='optimize')
jobs.append(graph)
# Train a graph on a device
def train(device, graph):
print "Start training on %s" % device
with tf.Session(graph=graph) as session:
x = graph.get_tensor_by_name('x:0')
y = graph.get_tensor_by_name('y:0')
cost = graph.get_tensor_by_name('cost:0')
optimize = graph.get_operation_by_name('optimize')
session.run(tf.global_variables_initializer())
batch_size = 500
for epoch in range(5):
total_batch = int(train_x_all.shape[0] / batch_size)
for i in range(total_batch):
batch_x = train_x_all[i * batch_size:(i + 1) * batch_size]
batch_y = train_y_all[i * batch_size:(i + 1) * batch_size]
_, c = session.run([optimize, cost], feed_dict={x: batch_x, y: batch_y})
if i % 20 == 0:
print "Device %s: epoch #%d step=%d cost=%f" % (device, epoch, i, c)
# Start threads in parallel
train_threads = []
for i, graph in enumerate(jobs):
train_threads.append(threading.Thread(target=train, args=(devices[i], graph)))
for t in train_threads:
t.start()
for t in train_threads:
t.join()
Note that train function operates with tensors and operations from the graph in the context, i.e. each cost and optimize is different.
This produces the following output, which shows that two models are trained in parallel:
Start training on /gpu
Start training on /cpu
Device /cpu: epoch #0 step=0 cost=2.302585
Device /cpu: epoch #0 step=20 cost=1.788247
Device /cpu: epoch #0 step=40 cost=1.400490
Device /cpu: epoch #0 step=60 cost=1.271820
Device /gpu: epoch #0 step=0 cost=2.302585
Device /cpu: epoch #0 step=80 cost=1.128214
Device /gpu: epoch #0 step=20 cost=2.105802
Device /cpu: epoch #0 step=100 cost=0.927004
Device /cpu: epoch #1 step=0 cost=0.905336
Device /gpu: epoch #0 step=40 cost=1.908744
Device /cpu: epoch #1 step=20 cost=0.865687
Device /gpu: epoch #0 step=60 cost=1.808407
Device /cpu: epoch #1 step=40 cost=0.754765
Device /gpu: epoch #0 step=80 cost=1.676024
Device /cpu: epoch #1 step=60 cost=0.794201
Device /gpu: epoch #0 step=100 cost=1.513800
Device /gpu: epoch #1 step=0 cost=1.451422
Device /cpu: epoch #1 step=80 cost=0.786958
Device /gpu: epoch #1 step=20 cost=1.415125
Device /cpu: epoch #1 step=100 cost=0.643715
Device /cpu: epoch #2 step=0 cost=0.674683
Device /gpu: epoch #1 step=40 cost=1.273473
Device /cpu: epoch #2 step=20 cost=0.658424
Device /gpu: epoch #1 step=60 cost=1.300150
Device /cpu: epoch #2 step=40 cost=0.593681
Device /gpu: epoch #1 step=80 cost=1.242193
Device /cpu: epoch #2 step=60 cost=0.640543
Device /gpu: epoch #1 step=100 cost=1.105950
Device /gpu: epoch #2 step=0 cost=1.089900
Device /cpu: epoch #2 step=80 cost=0.664947
Device /gpu: epoch #2 step=20 cost=1.088389
Device /cpu: epoch #2 step=100 cost=0.535446
Device /cpu: epoch #3 step=0 cost=0.580295
Device /gpu: epoch #2 step=40 cost=0.983053
Device /cpu: epoch #3 step=20 cost=0.566510
Device /gpu: epoch #2 step=60 cost=1.044966
Device /cpu: epoch #3 step=40 cost=0.518787
Device /gpu: epoch #2 step=80 cost=1.025607
Device /cpu: epoch #3 step=60 cost=0.562461
Device /gpu: epoch #2 step=100 cost=0.897545
Device /gpu: epoch #3 step=0 cost=0.907381
Device /cpu: epoch #3 step=80 cost=0.600475
Device /gpu: epoch #3 step=20 cost=0.911914
Device /cpu: epoch #3 step=100 cost=0.477412
Device /cpu: epoch #4 step=0 cost=0.527233
Device /gpu: epoch #3 step=40 cost=0.827964
Device /cpu: epoch #4 step=20 cost=0.513356
Device /gpu: epoch #3 step=60 cost=0.897128
Device /cpu: epoch #4 step=40 cost=0.474257
Device /gpu: epoch #3 step=80 cost=0.898960
Device /cpu: epoch #4 step=60 cost=0.514083
Device /gpu: epoch #3 step=100 cost=0.774140
Device /gpu: epoch #4 step=0 cost=0.799004
Device /cpu: epoch #4 step=80 cost=0.559898
Device /gpu: epoch #4 step=20 cost=0.802869
Device /cpu: epoch #4 step=100 cost=0.440813
Device /gpu: epoch #4 step=40 cost=0.732562
Device /gpu: epoch #4 step=60 cost=0.801020
Device /gpu: epoch #4 step=80 cost=0.815830
Device /gpu: epoch #4 step=100 cost=0.692840
You can try it yourself with standard MNIST data.
It's not ideal if there are many hyperparameters to tune, but you should be able to make an outer loop that iterates over possible hyperparameter tuples, assigns a particular graph to a device and runs them as shown above.

Answer 2

For a problem like this I would typical use the multiprocessing library instead of threading because compared to training a network the overhead of multiprocessing is small but removes any GIL problems. I think that this is the primary problem with your code. You are setting the "CUDA_VISIBLE_DEVICES" environment variable for each thread, but each thread is still sharing the same environment because they are in the same process.
So what I normally do in Tensorflow==2.1 is pass a GPU id number to the worker process which can then run the following code to set the visible GPU
gpus = tf.config.experimental.list_physical_devices('GPU')
my_gpu = gpus[gpu_id]
tf.config.set_visible_devices(my_gpu, 'GPU')
Tensorflow in that process will now only run on that one GPU
Sometimes the network you are training is small enough that you could actually run several at once on one GPU. To make sure that several can fit onto the GPU memory you can set the memory limit for each worker you start.
tf.config.set_logical_device_configuration(
my_gpu,
[tf.config.LogicalDeviceConfiguration(memory_limit=6000)]
)
But if you set the memory limit keep in mind that Tensorflow uses some extra memory outside of the limit for cuDNN or something so you need to have a bit of buffer memory as well for each session you run. Usually I just do trial an error to see what I can fit so sorry I don't have better numbers.