0

跨多个 GPU 的简单模型分区会导致工作负载在向前和向后传递期间从 GPU 移动到 GPU。任何时候，一个 GPU 都很忙。这是天真的版本。

with tf.device('/gpu:0'):
    model.add(Conv2D(32, kernel_size=(3, 3),
                 activation='relu',
                 input_shape=input_shape))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))

with tf.device('/gpu:1'):
    model.add(Conv2D(128, kernel_size=(3, 3),
                 activation='relu',
                 input_shape=input_shape))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(MaxPooling2D(pool_size=(2, 2)))

with tf.device('/gpu:2'):
    model.add(Dropout(0.25))
    model.add(Flatten())
    model.add(Dense(1500, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(num_classes, activation='softmax'))

我们需要示例代码（模板）来处理工作并通过发送批次波和协调每个 GPU 上的工作（前向、梯度计算、参数更新）来保持所有 GPU 忙碌。

这里通过使用 a提供了一个提示，data_flow_ops.StagingArea但一个具体的例子会有所帮助。

我知道数据分区（或数据并行）是可行的方法，但存在模型需要跨 CPU+GPU 分区的用例。

感谢任何指针或示例（伪）代码。

4

1 回答 1

0

模型并行性是一种在模型参数不适合单个 GPU 内存时使用的技术。当您的模型非常复杂（许多层）或某些层很大时，就会发生这种情况。通常模型并行是你应该只作为最后手段使用的东西，因为它通常很慢。

你的模型看起来很简单，所以我不确定你是否真的需要模型并行性（它只是一个例子吗？）。如果您一次只想使用单个 GPU，并且可以将所有模型放入单个 GPU，我不建议您进行模型并行。

如果您确定需要模型并行性，请参考此示例以使用 Apache MXNet 进行操作。

于 2018-08-07T19:42:02.650 回答