yolov8自己写train.py

yolo8自己写的train.py会运行报错,命令提示符里面就可以运行,环境安装的pip install -e .

img

  • 这篇文章讲的很详细,请看:yolov5train.py注释
  • 除此之外, 这篇博客: yolov3详细解释(内附代码详细说明)中的 train.py 部分也许能够解决你的问题, 你可以仔细阅读以下内容或跳转源博客中阅读:
  • """
    Retrain the YOLO model for your own dataset.
    """
    
    import numpy as np
    import keras.backend as K
    from keras.layers import Input, Lambda
    from keras.models import Model
    from keras.optimizers import Adam
    from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
    
    from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss
    from yolo3.utils import get_random_data
    
    #主函数
    def _main():
        annotation_path = 'train.txt'
        log_dir = 'logs/000/'
        classes_path = 'model_data/voc_classes.txt'
        anchors_path = 'model_data/yolo_anchors.txt'
        #得到类名,类总数,锚框
        class_names = get_classes(classes_path)
        num_classes = len(class_names)
        anchors = get_anchors(anchors_path)
    
        input_shape = (416,416) # multiple of 32, hw
        
        #判断是否用tiny_body
        is_tiny_version = len(anchors)==6 # default setting
        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/yolo_weights.h5') # make sure you know what you freeze
        
        #设置日志路径
        logging = TensorBoard(log_dir=log_dir)
        #设置检查点,3回检查一次,按验证集损失对比 保留最号的
        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)
        #学习衰减率 monitor 监测值  factor 缩放学习率的值,学习率将以lr = lr*factor的形式被减少
        #当patience个epoch过去而模型性能不提升时,学习率减少的动作会被触发
        reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)
        #早停 min_delta增大或减小的阈值,只有大于这个部分才算作improvement 
        #能够容忍多少个epoch内都没有improvement  
        early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=10, verbose=1)
    
        #分给验证集的比例
        val_split = 0.1
        with open(annotation_path) as f:
            lines = f.readlines()
        np.random.seed(10101)
        np.random.shuffle(lines)
        np.random.seed(None)
        #分配训练集和验证集的数量
        num_val = int(len(lines)*val_split)
        num_train = len(lines) - num_val
    
        # Train with frozen layers first, to get a stable loss.
        # Adjust num epochs to your dataset. This step is enough to obtain a not bad model.
        #先用冰冻层进行训练,得到稳定的损失。
        #调整num epoch到您的数据集。这一步足以获得一个不错的模型。
        if True:
            #设置优化器 loss设置不太懂????????????????????????????????????/
            model.compile(optimizer=Adam(lr=1e-3), loss={
                # use custom yolo_loss Lambda layer.
                'yolo_loss': lambda y_true, y_pred: y_pred})
    
            batch_size = 32
            print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
            #设置训练器
            model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
                    steps_per_epoch=max(1, num_train//batch_size),
                    validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
                    validation_steps=max(1, num_val//batch_size),
                    epochs=50,
                    initial_epoch=0,
                    callbacks=[logging, checkpoint])
            model.save_weights(log_dir + 'trained_weights_stage_1.h5')
    
        # Unfreeze and continue training, to fine-tune.
        # Train longer if the result is not good.
        #解冻并继续训练,以进行微调。
        #如果效果不好,就多训练。
        if True:
            #将模型层全部解冻
            for i in range(len(model.layers)):
                model.layers[i].trainable = True
            model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change
            print('Unfreeze all of the layers.')
    
            batch_size = 32 # note that more GPU memory is required after unfreezing the body
            print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
            model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
                steps_per_epoch=max(1, num_train//batch_size),
                validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
                validation_steps=max(1, num_val//batch_size),
                epochs=100,
                initial_epoch=50,
                callbacks=[logging, checkpoint, reduce_lr, early_stopping])
            model.save_weights(log_dir + 'trained_weights_final.h5')
    
        # Further training if needed.
        #如有需要,可接受进一步训练。
    
    #在_main中应用,得到类名一个列表
    def get_classes(classes_path):
        '''loads the classes'''
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names
    
    #在_main中应用,得到锚框(9,2)维度 或者 (6,2)维度
    def get_anchors(anchors_path):
        '''loads the anchors from a file'''
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)
    
    #在_train 中应用
    def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
                weights_path='model_data/yolo_weights.h5'):
        '''create the training model'''
        #参数解释:load_pretrained 是否下载预训练,
        #freeze_body 是否冻结预训练的层 1代表冻结darknet 2代表冻结yolo_body除3个检测层的其他层
        K.clear_session() # get a new session
        #占位参数 image_input 在经过Input后 shape会变成四个维度 多加一个batch_size维度 (None,None,None,3)
        image_input = Input(shape=(None, None, 3))
        h, w = input_shape#(416,416)
        num_anchors = len(anchors)
    
        #得到占位 y_true
        y_true = [Input(shape=(h//{0:32, 1:16, 2:8}[l], w//{0:32, 1:16, 2:8}[l], \
            num_anchors//3, num_classes+5)) for l in range(3)]
        
        #建立模型本体
        model_body = yolo_body(image_input, num_anchors//3, num_classes)
        print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))
        #是否加载预训练权重
        if load_pretrained:
            model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
            print('Load weights {}.'.format(weights_path))
            #是否冻结 代表不训练这些层
            if freeze_body in [1, 2]:
                # Freeze darknet53 body or freeze all but 3 output layers.
                #冻结暗网53的身体或冻结除3个输出层以外的所有层。
                num = (185, len(model_body.layers)-3)[freeze_body-1]
                for i in range(num): model_body.layers[i].trainable = False
                print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))
        #设置损失函数 *的作用是进行迭代,将列表中的元素按第一维度迭代处出来,实际就是减少一个维度
        model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
            arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(
            [*model_body.output, *y_true])
        #将Model 中 [model_body.input, *y_true]设置为输入 model_loss设置为输出
        model = Model([model_body.input, *y_true], model_loss)
    
        return model
    
    #在_train 中应用 与create_model相似
    def create_tiny_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
                weights_path='model_data/tiny_yolo_weights.h5'):
        '''create the training model, for Tiny YOLOv3'''
        K.clear_session() # get a new session
        image_input = Input(shape=(None, None, 3))
        h, w = input_shape
        num_anchors = len(anchors)
    
        y_true = [Input(shape=(h//{0:32, 1:16}[l], w//{0:32, 1:16}[l], \
            num_anchors//2, num_classes+5)) for l in range(2)]
    
        model_body = tiny_yolo_body(image_input, num_anchors//2, num_classes)
        print('Create Tiny YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))
    
        if load_pretrained:
            model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
            print('Load weights {}.'.format(weights_path))
            if freeze_body in [1, 2]:
                # Freeze the darknet body or freeze all but 2 output layers.
                num = (20, len(model_body.layers)-2)[freeze_body-1]
                for i in range(num): model_body.layers[i].trainable = False
                print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))
    
        model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
            arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.7})(
            [*model_body.output, *y_true])
        model = Model([model_body.input, *y_true], model_loss)
    
        return model
    
    def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):
        '''data generator for fit_generator'''
        #读取所有的数据行
        n = len(annotation_lines)
        i = 0
        while True:
            image_data = []
            box_data = []
            #对batch_size中的每一行数据进行处理
            for b in range(batch_size):
                #洗牌 打乱
                if i==0:
                    np.random.shuffle(annotation_lines)
                #得到经过调整后的数据
                image, box = get_random_data(annotation_lines[i], input_shape, random=True)
                image_data.append(image)
                box_data.append(box)
                #确保i在n里 也保证i一直增大
                i = (i+1) % n
            image_data = np.array(image_data)
            box_data = np.array(box_data)
            #对输入进行预处理 得到y_true
            y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
            #yield作用是输出后继续循环
            yield [image_data, *y_true], np.zeros(batch_size)
    
    def data_generator_wrapper(annotation_lines, batch_size, input_shape, anchors, num_classes):
        n = len(annotation_lines)
        if n==0 or batch_size<=0: return None
        return data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes)
    
    if __name__ == '__main__':
        _main()