TensorFlow如何将构建与执行分开以实现数据并行性

时间:2018-09-13 08:52:34

标签: python tensorflow

在TensorFlow中,通常将模块封装在函数或类中,然后抽象所需的变量的创建以调用它们

net = slim.fully_connected(inputs=input, num_output=neurons ..)
net = tf.layers.conv2d(net, num_filters, filter_size ..)

这里将首先为每个操作创建权重和偏差,然后重新使用。

当我们想要实现数据并行性时,我们希望将变量创建并存储在CPU上,然后将其与数据一起发送到GPU中,如下图所示

enter image description here

cifar10_multi_gpu_train示例中,如果您在同一目录中检查tf.layers,您会发现他们不使用cifar10.py,您会发现他们使用了较低级别的操作。 fully_connectedconv2d,并在CPU上手动创建内核,权重和偏差。 如果我们想使用已经为方便在TensorFlow中使用而实现的复杂结构,这可能会非常麻烦。

我的问题是:我们可以使用高级模块抽象(来自slim / tf.layers以及其他抽象变量创建的方式)的方式,该方式将在CPU上创建变量,但操作将在GPU上执行?

1 个答案:

答案 0 :(得分:1)

编辑:

关于将变量固定到CPU,可以使用tf.device和设备功能来完成。在分布式环境中,您拥有tf.train.replica_device_setter,但是针对本地情况很容易执行类似的操作:

import tensorflow as tf

def my_device_placement(device, vars_device='/cpu:0'):
    # Ops to pin on the CPU
    VAR_TYPES = ['Variable', 'VariableV2', 'VarHandleOp']
    def device_function(op):
        return vars_device if op.type in VAR_TYPES else device
    return device_function

def conv2d_replica(input_, filters, kernel_size, name, device, is_first_replica):
    with tf.device(my_device_placement(device)):
        return tf.layers.conv2d(input_, filters, kernel_size, name=name, reuse=not is_first_replica)

inp = tf.placeholder(tf.float32, [None, 100, 100, 3])
lyr1 = conv2d_replica(inp, 5, [20, 20], 'Layer', '/gpu:0', True)
lyr2 = conv2d_replica(inp, 5, [20, 20], 'Layer', '/gpu:1', False)
print('Device of first replica:', lyr1.device)
print('Device of second replica:', lyr2.device)
print('Variable devices:')
for var in tf.trainable_variables():
    print(var.name, var.device)

输出:

Device of first replica: /gpu:0
Device of second replica: /gpu:1
Variable devices:
Layer/kernel:0 /cpu:0
Layer/bias:0 /cpu:0

应该在CPU上执行的操作由您决定。您可以查看python/training/device_setter.py中的STANDARD_PS_OPS,以了解TensorFlow认为是固定在参数服务器上的标准操作集(在这种情况下,它是本地的,但是想法很相似)。


通过tf.layers,您可以使用namereuse参数。使用reuse=True时,图层将使用先前创建的具有相同name的图层的权重。请注意,这意味着您第一次创建图层reuse应该是False

import tensorflow as tf

inp = tf.placeholder(tf.float32, [None, 100, 100, 3])
lyr1 = tf.layers.conv2d(inp, 5, [20, 20], name='Layer', reuse=False)
lyr2 = tf.layers.conv2d(inp, 5, [20, 20], name='Layer', reuse=True)

图:

Model 1

这里BiasAdd节点是该层的输出。权重在同一层中创建,并在第二层中重复使用。

请注意,这甚至可以在命名空间中使用(我不确定是否有此意图,因为我还没有找到关于它的明确文档):

import tensorflow as tf

inp = tf.placeholder(tf.float32, [None, 100, 100, 3])
with tf.name_scope('Replica1'):
    lyr1 = tf.layers.conv2d(inp, 5, [20, 20], name='Layer', reuse=False)
with tf.name_scope('Replica2'):
    lyr2 = tf.layers.conv2d(inp, 5, [20, 20], name='Layer', reuse=True)

图:

Model 2

注意:即使如今已基本弃用它,看来tf.slim也提供相同的功能。在这种情况下,变量范围也有一个reuse参数,然后是一个scope参数,因此它类似于:

import tensorflow as tf

inp = tf.placeholder(tf.float32, [None, 10])
with tf.variable_scope('Layer') as scope:
    lyr1 = tf.contrib.slim.fully_connected(inp, 5, reuse=False, scope=scope)
with tf.variable_scope('Layer') as scope:
    lyr2 = tf.contrib.slim.fully_connected(inp, 5, reuse=True, scope=scope)