简单的TensorFlow示例将模型的一个副本加载到可用的每个GPU上

Question

我正在考虑从不同的生产机器学习框架移植到TensorFlow。在我们当前的培训和推理系统中，我们将模型的副本加载到机器上的多个GPU中。

我现在想保持这种负载均衡的方式。在哪里可以找到一个将TF模型的一个副本加载到机器上可用的每个GPU上的简单示例？

Answer 1

以下是https://github.com/rafaljozefowicz/lm/blob/master/language_model.py#L21

的示例

将模型创建代码包装到_forward函数中，然后为每个GPU调用一次

    for i in range(hps.num_gpus):
        with tf.device(assign_to_gpu(i, ps_device)), tf.variable_scope(tf.get_variable_scope(),
                                                                       reuse=True if i > 0 else None):
            loss = self._forward(i, xs[i], ys[i], ws[i])
            losses += [loss]
            if mode == "train":
                cur_grads = self._backward(loss, summaries=(i == hps.num_gpus - 1))
                tower_grads += [cur_grads]

    self.loss = tf.add_n(losses) / len(losses)