关于将变量固定到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
STANDARD_PS_OPS
在里面
python/training/device_setter.py
与
tf.layers
,您可以使用
name
和
reuse
参数。什么时候?
reuse=True
名称
. 注意这意味着第一次创建图层时
重新使用
应该是
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)
图表:
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)
图表:
注意:尽管现在基本上不赞成,但似乎
tf.slim
也提供相同的功能。在这种情况下,有一个
重新使用
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)