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: [], [], [], [].

 

相关推荐
©️2020 CSDN 皮肤主题: 游动-白 设计师:白松林 返回首页