Author: Zongwei Zhou | 周纵苇
Weibo: @MrGiovanni
Email: zongweiz@asu.edu

---

References.

官方文档：multi_gpu_model
以及Google

0. 误区

目前Keras是支持了多个GPU同时训练网络，非常容易，但是靠以下这个代码是不行的。

os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

当你监视GPU的使用情况（nvidia-smi -l 1）的时候会发现，尽管GPU不空闲，实质上只有一个GPU在跑，其他的就是闲置的占用状态，也就是说，如果你的电脑里面有多张显卡，无论有没有上面的代码，Keras都会默认的去占用所有能检测到的GPU。这行代码在你只需要一个GPU的时候时候用的，也就是可以让Keras检测不到电脑里其他的GPU。假设你一共有三张显卡，每个显卡都是有自己的标号的（0, 1, 2），为了不影响别人的使用，你只用其中一个，比如用gpu=1的这张，那么

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

然后再监视GPU的使用情况（nvidia-smi -l 1），确实只有一个被占用，其他都是空闲状态。所以这是一个Keras使用多显卡的误区，它并不能同时利用多个GPU。

1. 目的

为什么要同时用多个GPU来训练？

单个显卡内存太小 -> batch size无法设的比较大，有时甚至batch_size=1都内存溢出（OUT OF MEMORY）

从我跑深度网络的经验来看，batch_size设的大一点会比较好，相当于每次反向传播更新权重，网络都可以看到更多的样本，从而不会每次iteration都过拟合到不同的地方去Don't Decay the Learning Rate, Increase the Batch Size。当然，我也看过有论文说也不能设的过大，原因不明... 反正我也没有机会试过。我建议的batch_size大概就是64~256的范围内，都没什么大问题。

但是随着现在网络的深度越来越深，对于GPU的内存要求也越来越大，很多入门的新人最大的问题往往不是代码，而是从Github里面抄下来的代码自己的GPU太渣，实现不了，只能降低batch_size，最后训练不出那种效果。

解决方案两个：一是买一个超级牛逼的GPU，内存巨大无比；二是买多个一般般的GPU，一起用。

第一个方案不行，因为目前即便最好的NVIDIA显卡，内存也不过十几个G了不起了，网络一深也挂，并且买一个牛逼显卡的性价比不高。所以、学会在Keras下用多个GPU是比较靠谱的选择。

2. 实现

2.1 设计一个类

cite: parallel_model.py

import tensorflow as tf
import keras.backend as K
import keras.layers as KL
import keras.models as KM


class ParallelModel(KM.Model):
    """Subclasses the standard Keras Model and adds multi-GPU support.
    It works by creating a copy of the model on each GPU. Then it slices
    the inputs and sends a slice to each copy of the model, and then
    merges the outputs together and applies the loss on the combined
    outputs.
    """

    def __init__(self, keras_model, gpu_count):
        """Class constructor.
        keras_model: The Keras model to parallelize
        gpu_count: Number of GPUs. Must be > 1
        """
        self.inner_model = keras_model
        self.gpu_count = gpu_count
        merged_outputs = self.make_parallel()
        super(ParallelModel, self).__init__(inputs=self.inner_model.inputs,
                                            outputs=merged_outputs)

    def __getattribute__(self, attrname):
        """Redirect loading and saving methods to the inner model. That's where
        the weights are stored."""
        if 'load' in attrname or 'save' in attrname:
            return getattr(self.inner_model, attrname)
        return super(ParallelModel, self).__getattribute__(attrname)

    def summary(self, *args, **kwargs):
        """Override summary() to display summaries of both, the wrapper
        and inner models."""
        super(ParallelModel, self).summary(*args, **kwargs)
        self.inner_model.summary(*args, **kwargs)

    def make_parallel(self):
        """Creates a new wrapper model that consists of multiple replicas of
        the original model placed on different GPUs.
        """
        # Slice inputs. Slice inputs on the CPU to avoid sending a copy
        # of the full inputs to all GPUs. Saves on bandwidth and memory.
        input_slices = {name: tf.split(x, self.gpu_count)
                        for name, x in zip(self.inner_model.input_names,
                                           self.inner_model.inputs)}

        output_names = self.inner_model.output_names
        outputs_all = []
        for i in range(len(self.inner_model.outputs)):
            outputs_all.append([])

        # Run the model call() on each GPU to place the ops there
        for i in range(self.gpu_count):
            with tf.device('/gpu:%d' % i):
                with tf.name_scope('tower_%d' % i):
                    # Run a slice of inputs through this replica
                    zipped_inputs = zip(self.inner_model.input_names,
                                        self.inner_model.inputs)
                    inputs = [
                        KL.Lambda(lambda s: input_slices[name][i],
                                  output_shape=lambda s: (None,) + s[1:])(tensor)
                        for name, tensor in zipped_inputs]
                    # Create the model replica and get the outputs
                    outputs = self.inner_model(inputs)
                    if not isinstance(outputs, list):
                        outputs = [outputs]
                    # Save the outputs for merging back together later
                    for l, o in enumerate(outputs):
                        outputs_all[l].append(o)

        # Merge outputs on CPU
        with tf.device('/cpu:0'):
            merged = []
            for outputs, name in zip(outputs_all, output_names):
                # If outputs are numbers without dimensions, add a batch dim.
                def add_dim(tensor):
                    """Add a dimension to tensors that don't have any."""
                    if K.int_shape(tensor) == ():
                        return KL.Lambda(lambda t: K.reshape(t, [1, 1]))(tensor)
                    return tensor
                outputs = list(map(add_dim, outputs))

                # Concatenate
                merged.append(KL.Concatenate(axis=0, name=name)(outputs))
        return merged

2.2 调用非常简洁

GPU_COUNT = 3 # 同时使用3个GPU
model = keras.applications.densenet.DenseNet201() # 比如使用DenseNet-201
model = ParallelModel(model, GPU_COUNT)
model.compile(optimizer=Adam(lr=1e-5), loss='binary_crossentropy', metrics = ['accuracy'])
model.fit(X_train, y_train,
              batch_size=batch_size*GPU_COUNT, 
              epochs=nb_epoch, verbose=0, shuffle=True,
              validation_data=(X_valid, y_valid))

model.save_weights('/path/to/save/model.h5')