Question

我正在尝试计算跨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设备。

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

Answer 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，请尝试将渐变与相应的操作相对应。

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

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

1 个答案: