Question

在多个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繁忙。

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

我知道可以进行数据分区（或数据并行），但是在某些用例中，需要在CPU + GPU之间对模型进行分区。

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

Answer 1

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

您的模型看起来很简单，所以我不确定您是否真的需要模型并行性（仅仅是一个例子？）。如果您想一次只使用一个GPU，并且可以将所有模型都集成到一个GPU中，那么我不建议您进行模型并行化。

如果确定需要模型并行性，请参考this example来使用Apache MXNet。

TensorFlow中用于模型并行和流水线训练的示例代码

1 个答案: