yolov3的多显卡训练单显卡测试（keras）版本

最新推荐文章于 2023-08-02 12:45:36 发布

原创

最新推荐文章于 2023-08-02 12:45:36 发布 · 7.3k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#yolo

这篇博客记录了如何使用Keras进行Yolov3模型的多GPU训练，并在单GPU上进行测试。作者参考了多个资源，包括原始的keras-yolo3仓库、多GPU训练的实现以及关于`multi_gpu_model()`函数的详细解释。经过修改，实现了训练过程的多GPU支持，并通过CPU保存权重以便在单GPU环境下进行模型测试。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

今天是记录yolov3的多显卡的训练，之前也尝试过，没有跑起来，或者是跑起来了，但是不能单显卡检测，之后就不了了之了。

知道最近要训练大批量的数据，有多显卡不用实属浪费，于是就在网络找了找，跑成功了，可以多显卡训练，单显卡检测了。

主要是参考3个地方：

1.https://github.com/qqwweee/keras-yolo3 这个就是原始的yolo keras版本，但是好像没有多gpu选择

2.https://github.com/FlyEgle/keras-yolo3 这个就是多gpu训练的版本，然后在他的基础上修改了一下，他的代码只有训练前面的冻结层，我自己补全了一些东西，还是很感谢他分享代码的。

3.https://blog.youkuaiyun.com/u010122972/article/details/84784245 这个是讲multi_gpu_model（）这个函数的，巨详细，很好！

然后下面贴上我修改的FlyEgle大神的代码，也是一个记录。主要思想就是用多gpu训练，然后用单cpu保存权重，这样就可以在单gpu上test了（我也不知道为啥是用cpu保存权重。。。。。）

主要修改的是FlyEgle里train_height_point.py的代码

"""
Retrain the YOLO model for your own dataset.
"""
import os

import numpy as np
import json
import codecs
import matplotlib.pyplot as plt
import keras.backend as K
import tensorflow as tf

from keras.layers import Input, Lambda, add
from keras.models import Model
from keras.optimizers import Adam, SGD
from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping, LearningRateScheduler
from keras.backend.tensorflow_backend import set_session 

from yolo3.model import preprocess_true_boxes, yolo_body, yolo_loss
from yolo3.utils import get_random_data
# 当前程序可以调用的gpu
os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"

def _main():
    # 这里是要的由VOC生成的txt，里面包括image地址和框和种类 需修改
    annotation_path = '/data/drone_detect/keras-yolo3-multigpu/6.6/final_label.txt'
    # 是存放生成好的新的weights的
    log_dir = '/data/drone_detect/keras-yolo3-multigpu/6.6/weights/'
    # 这个我只标记了4类，所以要修改，原来voc是有20类的
    classes_path = '/data/drone_detect/keras-yolo3-multigpu/6.6/voc_classes.txt'
    # 这里的anchor暂时先不修改(最好从新kmean算anchors)
    anchors_path = '/data/drone_detect/keras-yolo3-multigpu/6.6/all_anchors.txt'
    class_names = get_classes(classes_path)
    num_classes = len(class_names)
    anchors = get_anchors(anchors_path)

    input_shape = (480, 960) # multiple of 32, hw
    is_tiny_version = len(anchors)==6 # default setting
    # use tiny model if need else darkent53 model for train
    if is_tiny_version:
        '''cpu filed for create model and save weights from model'''
        pass
        # with tf.device('/cpu:0'):
        #     template_model = create_tiny_model(input_shape, anchors, num_classes,
        #         freeze_body=1, weights_path='model_data/tiny_yolo_weights.h5')
        #     print('model.input shape', model.input_layers)
    else:
        '''build the darknet53 model on the cpu'''
        with tf.device('/cpu:0'):
            # 这里放预训练权重
            # 这里是创建训练模型，在这里还没有加入duogpu模式
            template_model = create_model(input_shape, anchors, num_classes,
                freeze_body=1, weights_path="/data/drone_detect/keras-yolo3-multigpu/6.5/trained_weights_final_480x960_30.62.h5") # make sure you know what you freeze
            print('model.input shape', template_model.input_layers, 'model.inputs', template_model.inputs)
    logging = TensorBoard(log_dir=log_dir)
    checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
                                 monitor='val_loss', save_weights_only=True, save_best_only=True, period=3)
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)

    # template_model input
    # print('model inputs', model.input)
    '''add shedule for reduce lr ratio by epoch monitor'''
    def shedule(epoch):
        if epoch <= 150:
            return 1e-3
        elif epoch <= 170:
            return 1e-4
        else:
            return 1e-5
    
    lr_shedule = LearningRateScheduler(shedule)
    
    early_stopping = EarlySto