multithreading - 使用 tensorflow 在并行 GPU 上运行超参数优化

Question

我有一个训练函数，可以在这里端到端地训练一个 tf 模型(仅为了说明而设计):

def opt_fx(params, gpu):
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu

    sess = tf.Session()
    # Run some training on a particular gpu...
    sess.run(...)

我想使用每个 gpu 的模型在 20 次试验中运行超参数优化:

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()

然而，这失败了……正确的方法是什么？

Answer 1

我不确定这是否是最好的方法，但我最终做的是为每个设备定义一个图表，并在单独的 session 中训练每个图表。这可以并行化。我试图在不同的设备中重复使用该图，但那没有用。这是我的版本在代码中的样子(一个完整的例子):

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()

请注意，train 函数在上下文中使用来自 graph 的张量和操作，即每个 cost 和 optimize 是不同的。

这会产生以下输出，表明两个模型是并行训练的:

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

您可以使用 standard MNIST data 自行尝试.

如果有许多超参数需要调整，这并不理想，但您应该能够制作一个外部循环，迭代可能的超参数元组，将特定图形分配给设备并如上所示运行它们。