Question

是否可以告诉tf.contrib.layers中定义的层存储定义CPU上的内核和偏差的变量（例如/cpu:0），同时仍在其他设备上执行像卷积这样的操作？< / p>

Answer 1

我似乎找到了一种解决方法，它允许我使用tf.contrib.layer中的图层保留我的模型描述。诀窍是在使用get_variable方法创建变量时，将图层中的变量节点添加到图形中。

例如：

from tensorflow.contrib.framework.python.ops import variables

layer_scope = 'fully_connected'
with tf.device('/cpu:0'):
    weights = variables.model_variable(layer_scope+'/weights', shape=...)
    biases = variables.model_variable(layer_scope+'/biases', shape=...)

fc = tf.contrib.layers.fully_connected(some_input, ...)

可以找到here at the tensorflow github函数model_variable。

记录设备放置并查看张量板中的图表现在可以看到变量放在cpu上，而可以将其他操作放在fully_connected上{ {1}}。

下面你可以找到我放在一起测试这种行为的脚本。在运行时，它将创建一个名为gpu的文件夹，其中可以找到tensorboard摘要和图表。

在我看来这非常hacky，但在我看来，你不会丢失任何功能，而不会增加许多代码行。

如果有人有兴趣对此进行测试，我会对反馈感到满意。

layer_device_placement_summary

编辑：

我在github上发出了显式变量放置的功能请求。你可以在这里讨论： issue9517

Answer 2

有一种更简单的方法：

def net(x):
    res = tf.contrib.layers.conv2d(x, 10, [3, 3])
    # and define other layers
    return res


with tf.device('/CPU:0'):
    with tf.name_scope('CPU'):
        net(x)
for i in range(num_gpus):
    with tf.device('/GPU:%d' % i):
        with tf.name_scope('GPU_%d' % i):
            with tf.variable_scope(tf.get_variable_scope(), reuse=True):
                y = net(x)
                loss = ...
                grad = ...

设备在使用tf.contrib.layers时放置内核和偏差

2 个答案:

编辑：