我正在尝试计算跨 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，则仅在 /gpu:1 中执行 tf.reduce_mean 操作。那么有没有什么简单的方法可以将梯度计算的操作自动分配给对应的GPU设备呢？
请您参考如下方法:
来自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:如果为真，则尝试将梯度与相应的操作并置。
https://github.com/tensorflow/tensorflow/issues/16328#issuecomment-359899310
