Tensorflow 行为：跨多 GPU 的梯度计算

Posted 2023-02-16

【中文标题】Tensorflow 行为：跨多 GPU 的梯度计算

【英文标题】：Tensorflow behavior: gradient computation across multi-GPU

【发布时间】：2018-01-21 17:11:00

【问题描述】：

我正在尝试计算跨 GPU 的批量归一化层的全局均值和全局方差，应该考虑前向和后向。 使用\sigma^2 = mean(x^2) - mean(x)^2，渐变 w.r.t.每个x 都可以在x 所附加的GPU 中独立计算。

但是，在计算梯度时，我遇到了一个问题：没有指定 GPU 设备，tf.gradient 将使用\gpu:0。 我无法指定梯度计算的每个操作，因为梯度是由optimizer 自动计算的，并且只计算参数的梯度。

我的问题是，如果一个节点显式附加到 GPU 设备，为什么梯度不能附加到同一个 GPU 设备？

我尝试了这段代码，得到了两个时间线文件timelines.zip 和下面的两个快照。

import tensorflow as tf
import numpy as np
from tensorflow.python.client import timeline

N_SAMPLES = 100000000


def all_reduce(gpu_num):
    means = []
    x2s = []
    axs = []
    for i in range(gpu_num):
        with tf.device('/cpu:0'):
            x = tf.placeholder(dtype=tf.float32, shape=[N_SAMPLES], name='local_input_%d' % i)
        with tf.device('/gpu:%d'%i):
            ax = tf.multiply(10.0, x, name='local_multiply_%d'%i)
            mean = tf.reduce_mean(ax, name='local_mean_%d'%i)
            x2 = tf.square(ax, name='local_square_%d'%i)
            axs.append(ax)
            means.append(mean)
            x2s.append(x2)

    with tf.device('/gpu:0'):
        global_mean = tf.reduce_mean(means, name='global_mean')
        global_var = tf.subtract(tf.reduce_mean(x2s, name='global_x2'),
                                 tf.square(global_mean, name='global_mean_square'),
                                 name='global_sub')
        print global_var.get_shape()

    gs = []
    # manually
    # for i in range(gpu_num):
    #     with tf.device('/gpu:%d'%i):
    #         gradient_wrt_mean = tf.gradients(global_mean, axs[i])
    #         gradient_wrt_var = tf.gradients(global_var, axs[i])
    #         gs.append(gradient_wrt_mean)
    #         gs.append(gradient_wrt_var)

    # auto by tf
    gradient_wrt_mean = tf.gradients(global_mean, axs)
    gradient_wrt_var = tf.gradients(global_var, axs)
    gs.append(gradient_wrt_var)
    gs.append(gradient_wrt_mean)

    for n in tf.get_default_graph().as_graph_def().node:
        print [n.name, n.device]

    return global_mean, global_var, axs, gs


def main(_):
    gpu_num = 2
    mean_op, var_op, xs, gs = all_reduce(gpu_num)
    x = np.random.randn(N_SAMPLES*gpu_num)
    print np.mean(x), np.var(x)
    feed_dict = dict()
    for i in range(gpu_num):
        feed_dict[xs[i]] = x[i*N_SAMPLES:(i+1)*N_SAMPLES]

    run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
    run_metadata = tf.RunMetadata()
    gpu_options = tf.GPUOptions(allow_growth=False)
    config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options)
    sess = tf.Session(config=config)

    # mean, var, g = sess.run([
    #     mean_op, var_op, gs
    # ], feed_dict=feed_dict, options=run_options, run_metadata=run_metadata)
    # print mean, var

    g = sess.run([
        gs
    ], feed_dict=feed_dict, options=run_options, run_metadata=run_metadata)

    # Create the Timeline object, and write it to a json
    tl = timeline.Timeline(run_metadata.step_stats)
    ctf = tl.generate_chrome_trace_format()
    with open('timeline.json', 'w') as f:
        f.write(ctf)


if __name__ == '__main__':
    tf.app.run()

两个数字： 自动，不指定 GPU 设备。

手动指定 GPU 设备。

如果使用 tf.gradient 而不指定 GPU 设备，则在 /gpu:1 中仅执行 tf.reduce_mean 操作。那么是否有一些简单的方法可以将梯度计算的操作自动分配给对应的 GPU 设备？

【参考方案1】：

从github回答：

tf.gradients(
ys,
xs,
grad_ys=None,
name='gradients',
colocate_gradients_with_ops=False,
gate_gradients=False,
aggregation_method=None,
stop_gradients=None
)

colocate_gradients_with_ops：如果为 True，请尝试将渐变与相应的 op 共置。

https://github.com/tensorflow/tensorflow/issues/16328#issuecomment-359899310