如何在 TensorFlow 中加速具有多个 GPU 的 RNN 模型的训练？答案

【问题标题】：How to speed up the training of an RNN model with multiple GPUs in TensorFlow?如何在 TensorFlow 中加速具有多个 GPU 的 RNN 模型的训练？
【发布时间】：2018-05-25 13:33:28
【问题描述】：

例如，RNN 是一个动态的 3 层双向 LSTM，隐藏向量大小为 200 (tf.nn.bidirectional_dynamic_rnn)，我有 4 个 GPU 来训练模型。我看到post 在一批样本子集上使用data parallelism，但这并没有加快训练过程。

【问题讨论】：

标签： tensorflow distributed-computing lstm recurrent-neural-network multiple-gpu

【解决方案1】：

你也可以试试model parallelism。一种方法是制作这样的单元格包装器，它将在特定设备上创建单元格：

class DeviceCellWrapper(tf.nn.rnn_cell.RNNCell):
  def __init__(self, cell, device):
    self._cell = cell
    self._device = device

  @property
  def state_size(self):
    return self._cell.state_size

  @property
  def output_size(self):
    return self._cell.output_size

  def __call__(self, inputs, state, scope=None):
    with tf.device(self._device):
      return self._cell(inputs, state, scope)

然后将每个单独的层放到专用 GPU 上：

cell_fw = DeviceCellWrapper(cell=tf.nn.rnn_cell.LSTMCell(num_units=n_neurons, state_is_tuple=False), device='/gpu:0')
cell_bw = DeviceCellWrapper(cell=tf.nn.rnn_cell.LSTMCell(num_units=n_neurons, state_is_tuple=False), device='/gpu:0')
outputs, states = tf.nn.bidirectional_dynamic_rnn(cell_fw, cell_bw, X, dtype=tf.float32)

【讨论】：