Как ускорить обучение модели RNN с несколькими графическими процессорами в TensorFlow?

Question

Например, RNN является динамическим 3-уровневым двунаправленным LSTM с размером скрытого вектора 200 (tf.nn.bidirectional_dynamic_rnn) и у меня есть 4 графических процессора для обучения модели. Я видел пост, используя data parallelism на подмножествах образцов в партии, но это не ускорило процесс обучения.

Answer 1

Вы также можете попробовать модель параллелизма. Один из способов сделать это - создать такую ​​обертку для ячеек, которая будет создавать ячейки на определенном устройстве:

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)

Затем поместите каждый отдельный слой на выделенный графический процессор:

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)