我正在尝试为跨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设备，则在tf.reduce_mean中仅执行/gpu:1操作。那么，是否有一些简单的方法可以将梯度计算的操作自动分配给相应的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：如果为True，请尝试将渐变与相应的op并置。
https://github.com/tensorflow/tensorflow/issues/16328#issuecomment-359899310