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

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 来实现。