如何将多层 RNN 的每一层固定到 TensorFlow 中的不同 GPU？答案

【问题标题】：How do you pin each layer of a multi-layer RNN to a different GPU in TensorFlow?如何将多层 RNN 的每一层固定到 TensorFlow 中的不同 GPU？
【发布时间】：2016-09-18 11:12:34
【问题描述】：

我有一个由 LSTM 单元组成的多层 RNN。我想将每一层固定到不同的 GPU。你如何在 TensorFlow 中做到这一点？

import tensorflow as tf

n_inputs = 5
n_outputs = 100
n_layers = 5
n_steps = 20

X = tf.placeholder(tf.float32, shape=[None, n_steps, n_inputs])
lstm_cell = tf.nn.rnn_cell.BasicLSTMCell(num_units=n_outputs, state_is_tuple=True)
multi_layer_cell = tf.nn.rnn_cell.MultiRNNCell([lstm_cell]*n_layers, state_is_tuple=True)
outputs, states = tf.nn.dynamic_rnn(multi_layer_cell, X, dtype=tf.float32)

【问题讨论】：

标签： tensorflow gpu recurrent-neural-network lstm

【解决方案1】：

经过一番谷歌搜索后，我找到了 Nicolas Ivanov 的 this code。诀窍是通过扩展 RNNCell 抽象类来创建自己的单元格包装类。

这是我的看法：

import tensorflow as tf

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)

然后您可以像使用所有其他包装器一样使用此包装器：

n_inputs = 5
n_outputs = 100
devices = ["/gpu:0", "/gpu:1", "/gpu:2", "/gpu:3", "/gpu:4"]
n_steps = 20
X = tf.placeholder(tf.float32, shape=[None, n_steps, n_inputs])
lstm_cells = [DeviceCellWrapper(device, tf.nn.rnn_cell.BasicLSTMCell(
                                num_units=n_outputs, state_is_tuple=True))
              for device in devices]
multi_layer_cell = tf.nn.rnn_cell.MultiRNNCell(lstm_cells, state_is_tuple=True)
outputs, states = tf.nn.dynamic_rnn(multi_layer_cell, X, dtype=tf.float32)

【讨论】：

【解决方案2】：

我们通常看到两种不同的方法来做到这一点：或者像 MiniQuark 指出的那样环绕 BasicLSTMCell，或者使用不同的 MultiRNNCell 实现。环绕 BasicLSTMCell 可能是您的用例的更好选择。

【讨论】：