tensorflow - 使用 tf.device() 在 TensorFlow 中不起作用

Question

我有两个 GPU，想在 TensorFlow 中尝试一些分布式训练(模型并行)。

这两个 GPU 是:

/job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: TITAN Xp COLLECTORS EDITION, pci bus id: 0000:04:00.0, compute capability: 6.1
/job:localhost/replica:0/task:0/device:GPU:1 -> device: 1, name: TITAN X (Pascal), pci bus id: 0000:82:00.0, compute capability: 6.1

我的计划是将 LeNet 分成两部分，将每一部分分配给一个 GPU。

LeNet 有 5 层，我使用 with tf.device('/gpu:0'): 将第 1 层分配给 GPU 0，with tf.device('/gpu: 1'):将 layer2-layer5 分配给 GPU 1。

我知道没有必要在这个模型中进行模型并行化，但我只想在小模型中尝试模型并行化。

设备映射的日志显示，所有操作都已按我的意愿分配给设备:

layer5/fc3_b: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:1
layer5/fc3_b/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
layer5/fc3_b/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:1
layer5/fc3_w: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:1
layer5/fc3_w/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
layer5/truncated_normal/TruncatedNormal: (TruncatedNormal): /job:localhost/replica:0/task:0/device:GPU:1
layer5/truncated_normal/mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:1
layer5/truncated_normal: (Add): /job:localhost/replica:0/task:0/device:GPU:1
layer5/fc3_w/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:1
layer4/fc2_b: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:1
layer4/fc2_b/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
layer4/fc2_b/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:1
layer4/fc2_w: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:1
layer4/fc2_w/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
layer4/truncated_normal/TruncatedNormal: (TruncatedNormal): /job:localhost/replica:0/task:0/device:GPU:1
layer4/truncated_normal/mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:1
layer4/truncated_normal: (Add): /job:localhost/replica:0/task:0/device:GPU:1
layer4/fc2_w/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:1
layer3/fc1_b: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:1
layer3/fc1_b/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
layer3/fc1_b/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:1
layer3/fc1_w: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:1
layer3/fc1_w/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
layer3/truncated_normal/TruncatedNormal: (TruncatedNormal): /job:localhost/replica:0/task:0/device:GPU:1
layer3/truncated_normal/mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:1
layer3/truncated_normal: (Add): /job:localhost/replica:0/task:0/device:GPU:1
layer3/fc1_w/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:1
layer2/conv2_b: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:1
layer2/conv2_b/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
layer2/conv2_b/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:1
layer2/conv2_w: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:1
layer2/conv2_w/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:1
layer2/truncated_normal/TruncatedNormal: (TruncatedNormal): /job:localhost/replica:0/task:0/device:GPU:1
layer2/truncated_normal/mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:1
layer2/truncated_normal: (Add): /job:localhost/replica:0/task:0/device:GPU:1
layer2/conv2_w/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:1
init/NoOp_1: (NoOp): /job:localhost/replica:0/task:0/device:GPU:1
layer1/conv1_b: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:0
layer1/conv1_b/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:0
layer1/conv1_b/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:0
layer1/conv1_w: (VariableV2): /job:localhost/replica:0/task:0/device:GPU:0
layer1/conv1_w/read: (Identity): /job:localhost/replica:0/task:0/device:GPU:0
layer1/truncated_normal/TruncatedNormal: (TruncatedNormal): /job:localhost/replica:0/task:0/device:GPU:0
layer1/truncated_normal/mul: (Mul): /job:localhost/replica:0/task:0/device:GPU:0
layer1/truncated_normal: (Add): /job:localhost/replica:0/task:0/device:GPU:0
layer1/conv1_w/Assign: (Assign): /job:localhost/replica:0/task:0/device:GPU:0

但是我在timeline.json中得到了不同的结果，如下图所示。

时间线显示，好像layer2-layer5的ops是在GPU1上启动的，但是在GPU0上运行的。通过使用 with tf.device('/gpu:1'):，我认为这不是我想要的。

这在 TensorFlow 中是预期的吗？

这是我第一次问关于stack overflow的问题，如果需要其他信息，请告诉我，谢谢。

Answer 1

这只是 Chrome 跟踪事件格式的产物。

Stream, "/job:localhost/replica:0/task:0/devuce:GPU:0 Compute" 显示了为在 GPU:0 上执行的操作启动/排队 CUDA 内核的时间.

Stream, "/job:localhost/replica:0/task:0/devuce:GPU:1 Compute" 显示为在 GPU:1 上执行的操作启动/排队 CUDA 内核的时间.

所有流匹配，"/device:GPU:0/stream.* Compute"显示在所有 GPU 上实际执行操作的时间。要找出实际在哪个 GPU 上执行操作，您需要查看流 "/job:localhost/replica:0/task:0/devuce:GPU:.* Compute"

希望这能回答你的问题