keras之多GPU训练方法

在上一篇博客中，我们利用keras框架训练yolov3，训练脚本默认采用的是一块GPU，由于我们有多块GPU，因此可以设置多块GPU训练来加快训练速度。

实现方法很简单，首先在头文件中添加以下内容

from keras.utils import multi_gpu_model

然后找到自己构建网络的地方，在我这里，第一次构建model是以下的语句

if is_tiny_version:
    model = create_tiny_model(input_shape, anchors, num_classes,
        freeze_body=2, weights_path='model_data/tiny_yolo_weights.h5')
else:
    model = create_model(input_shape, anchors, num_classes,
        freeze_body=2, weights_path='model_data/trained_weights_80.h5') # make sure you know what you freeze

然后就找到create_tiny_model和create_model函数，进入函数内部，找到构建model的语句，比如我找create_model()函数，找到以下语句

model_body = yolo_body(image_input, num_anchors//3, num_classes)
print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))

然后在后边添加

model_body = multi_gpu_model(model_body,gpus=2)

表示使用两块GPU

如果想指定使用哪两块GPU，可以在开头添加如下语句

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"

表示使用第3第4块GPU

注意multi_gpu_model这个函数一定要放对位置，否则会报如下的错

tensorflow.python.framework.errors_impl.InvalidArgumentError: Can't concatenate scalars (use tf.stack instead) for 'yolo_loss_1/concat' (op: 'ConcatV2') with input shapes: [], [], [], [].