Question

我使用估算器和MirroredStrategy构建模型。它可以解决一些性能问题。我注意到所有变量都是MirroredVariable并放置在GPU上。我更喜欢选择一部分变量并将其放在cpu上，因此我将这些变量移至“ with tf.device（'/ cpu：0'）：”作用域下。但这是行不通的。当我检查设备放置日志时。这些变量仍在GPU上。我该如何解决？

我的一些代码：

mirrored_strategy = tf.contrib.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1","/gpu:2","/gpu:3"])
estimator_config = tf.estimator.RunConfig(train_distribute=mirrored_strategy,session_config=tf.ConfigProto(log_device_placement=True))

def model_fn():
    #other parts of my model with most variables
    with tf.device('/cpu:0'):
        var = tf.get_variable('weights', shape, initializer=initializer, dtype=dtype, trainable=trainable)
    #other model codes

设备放置日志：

conv1/weights/Initializer/truncated_normal/TruncatedNormal: (TruncatedNormal): /job:localhost/replica:0/task:0/device:GPU:0
conv1/weights/Initializer/truncated_normal/mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:0
conv1/weights/Initializer/truncated_normal: (Add): /job:localhost/replica:0/task:0/device:GPU:0
conv1/weights/IsInitialized/VarIsInitializedOp: (VarIsInitializedOp): /job:localhost/replica:0/task:0/device:GPU:0
conv1/weights/Assign: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
conv1/weights/Read/ReadVariableOp: (ReadVariableOp): /job:localhost/replica:0/task:0/device:GPU:0
conv1/weights/replica_1/Initializer/Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
conv1/weights/replica_1/IsInitialized/VarIsInitializedOp: (VarIsInitializedOp): /job:localhost/replica:0/task:0/device:GPU:1
conv1/weights/replica_1/Assign: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:1
conv1/weights/replica_1/Read/ReadVariableOp: (ReadVariableOp): /job:localhost/replica:0/task:0/device:GPU:1
conv1/weights/replica_2/Initializer/Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:2
conv1/weights/replica_2/IsInitialized/VarIsInitializedOp: (VarIsInitializedOp): /job:localhost/replica:0/task:0/device:GPU:2
conv1/weights/replica_2/Assign: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:2
conv1/weights/replica_2/Read/ReadVariableOp: (ReadVariableOp): /job:localhost/replica:0/task:0/device:GPU:2
conv1/weights/replica_3/Initializer/Identity: (Identity): /job:localhost/replica:0/task:0/device:GPU:3
conv1/weights/replica_3/IsInitialized/VarIsInitializedOp: (VarIsInitializedOp): /job:localhost/replica:0/task:0/device:GPU:3
conv1/weights/replica_3/Assign: (AssignVariableOp): /job:localhost/replica:0/task:0/device:GPU:3
conv1/weights/replica_3/Read/ReadVariableOp: (ReadVariableOp): /job:localhost/replica:0/task:0/device:GPU:3

使用MirroredStrategy时如何在CPU上放置变量？

0 个答案: