pytorch版本用自己数据集复现Yolov4检测
代码准备以及环境的搭建
yolov4-pytorch
这是YOLOv4架构的PyTorch重新实现，它基于官方的Darknet实现AlexeyAB / darknet与PASCAL VOC，COCO数据集，是众多pytorch复现版本中的一个。
本文所用复现环境：
windows+python3.6+pytorch1.2+apex0.1，其余见代码中requirements.txt
下载代码到指定文件夹：
在conda环境中，切换到该文件夹根目录,执行：
1：pip3 install -r requirements.txt --user
2：到官网安装pytorch
3：安装apex
https://github.com/NVIDIA/apex/
（1）在github上下载apex文件到本地
https://github.com/NVIDIA/apex
（2）然后将下载好的apex文件放到anaconda环境下的文件夹中。
anaconda3/envs/环境名/Lib/site-packages
（3)在conda,虚拟环境命令窗口，进入到apex文件夹中，激活环境
(4) 运行：python setup.py install
至此不常用的包都安装完了，剩下的pip即可
制作数据集
VOC的格式数据集过于常用，这里数据集制作基于voc格式，先布置好下面格式。
./data
└── VOCdevkit
└── VOC2007
├── Annotations # 标注的VOC格式的xml标签文件
├── JPEGImages # 数据集图片
├── ImageSet
│ └── Main
│ ├── test.txt # 划分的测试集
│ ├── train.txt # 划分的训练集
│ ├── trainval.txt
│ └── val.txt # 划分的验证集
先修改config/yolov4_config.py ，再运行utils/voc.py该代码，源代码中含有voc2012的下面代码以及注释好了。
# coding=utf-8
# project
DATA_PATH = r"D:\research\pytorch-YOLOv4-master\data"##voc数据集所在路径
PROJECT_PATH = "D:/research/xin-YOLOv4-pytorch-master/data"
DETECTION_PATH = "D:/research/xin-YOLOv4-pytorch-master/"
MODEL_TYPE = {"TYPE": 'YOLOv4'} #YOLO type:YOLOv4, Mobilenet-YOLOv4 or Mobilenetv3-YOLOv4
CONV_TYPE = {"TYPE": 'DO_CONV'} #conv type:DO_CONV or GENERAL
ATTENTION = {"TYPE": 'NONE'} #attention type:SEnet、CBAM or NONE
# train
TRAIN = {
"DATA_TYPE": 'VOC', #DATA_TYPE: VOC ,COCO or Customer设置voc类型的数据格式
"TRAIN_IMG_SIZE": 800,
"AUGMENT": True,
"BATCH_SIZE": 2,
"MULTI_SCALE_TRAIN": False,
"IOU_THRESHOLD_LOSS": 0.5,
"YOLO_EPOCHS": 50,
"Mobilenet_YOLO_EPOCHS": 120,
"NUMBER_WORKERS": 0,
"MOMENTUM": 0.9,
"WEIGHT_DECAY": 0.0005,
"LR_INIT": 1e-4,
"LR_END": 1e-6,
"WARMUP_EPOCHS": 2 # or None
}
# val
VAL = {
"TEST_IMG_SIZE": 800,
"BATCH_SIZE": 2,
"NUMBER_WORKERS": 0,
"CONF_THRESH": 0.005,
"NMS_THRESH": 0.45,
"MULTI_SCALE_VAL": True,
"FLIP_VAL": True,
"Visual": True
}
Customer_DATA = {"NUM": 1, #your dataset number
"CLASSES":['lf'],# your dataset class
}
'''
VOC_DATA = {"NUM": 20, "CLASSES":['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant', 'sheep', 'sofa',
'train', 'tvmonitor'],
}
'''
VOC_DATA = {"NUM": 1, "CLASSES":['lf'],#修改自己的类别，
}
COCO_DATA = {"NUM": 1, "CLASSES":['lf'],
}
'''
COCO_DATA = {"NUM":80,"CLASSES":['person',
'bicycle',
'car',
'motorcycle',
'airplane',
'bus',
'train',
'truck',
'boat',
'traffic light',
'fire hydrant',
'stop sign',
'parking meter',
'bench',
'bird',
'cat',
'dog',
'horse',
'sheep',
'cow',
'elephant',
'bear',
'zebra',
'giraffe',
'backpack',
'umbrella',
'handbag',
'tie',
'suitcase',
'frisbee',
'skis',
'snowboard',
'sports ball',
'kite',
'baseball bat',
'baseball glove',
'skateboard',
'surfboard',
'tennis racket',
'bottle',
'wine glass',
'cup',
'fork',
'knife',
'spoon',
'bowl',
'banana',
'apple',
'sandwich',
'orange',
'broccoli',
'carrot',
'hot dog',
'pizza',
'donut',
'cake',
'chair',
'couch',
'potted plant',
'bed',
'dining table',
'toilet',
'tv',
'laptop',
'mouse',
'remote',
'keyboard',
'cell phone',
'microwave',
'oven',
'toaster',
'sink',
'refrigerator',
'book',
'clock',
'vase',
'scissors',
'teddy bear',
'hair drier',
'toothbrush',]}
'''
# model
MODEL = {"ANCHORS":[[(1.25, 1.625), (2.0, 3.75), (4.125, 2.875)], # Anchors for small obj(12,16),(19,36),(40,28)
[(1.875, 3.8125), (3.875, 2.8125), (3.6875, 7.4375)], # Anchors for medium obj(36,75),(76,55),(72,146)
[(3.625, 2.8125), (4.875, 6.1875), (11.65625, 10.1875)]], # Anchors for big obj(142,110),(192,243),(459,401)
"STRIDES":[8, 16, 32],
"ANCHORS_PER_SCLAE":3
}
运行下面文件，会得到两个yolov4需要的数据集格式。
import sys
sys.path.append("..")
import xml.etree.ElementTree as ET
import config.yolov4_config as cfg
import os
from tqdm import tqdm
def parse_voc_annotation(data_path, file_type, anno_path, use_difficult_bbox=False):
"""
phase pascal voc annotation, eg:[image_global_path xmin,ymin,xmax,ymax,cls_id]
:param data_path: eg: VOC\VOCtrainval-2007\VOCdevkit\VOC2007
:param file_type: eg: 'trainval''train''val'
:param anno_path: path to ann file
:param use_difficult_bbox: whither use different sample
:return: batch size of data set
"""
classes = cfg.VOC_DATA["CLASSES"]
img_inds_file = os.path.join(data_path, 'ImageSets', 'Main', file_type+'.txt')
with open(img_inds_file, 'r') as f:
lines = f.readlines()
image_ids = [line.strip() for line in lines]
with open(anno_path, 'a') as f:
for image_id in tqdm(image_ids):
image_path = os.path.join(data_path, 'JPEGImages', image_id + '.bmp')
annotation = image_path
label_path = os.path.join(data_path, 'Annotations', image_id + '.xml')
root = ET.parse(label_path).getroot()
objects = root.findall('object')
for obj in objects:
difficult = obj.find("difficult").text.strip()
if (not use_difficult_bbox) and (int(difficult) == 1): # difficult 表示是否容易识别，0表示容易，1表示困难
continue
bbox = obj.find('bndbox')
class_id = classes.index(obj.find("name").text.lower().strip())
xmin = bbox.find('xmin').text.strip()
ymin = bbox.find('ymin').text.strip()
xmax = bbox.find('xmax').text.strip()
ymax = bbox.find('ymax').text.strip()
annotation += ' ' + ','.join([xmin, ymin, xmax, ymax, str(class_id)])
annotation += '\n'
# print(annotation)
f.write(annotation)
return len(image_ids)
if __name__ =="__main__":
# train_set : VOC2007_trainval 和 VOC2012_trainval
train_data_path_2007 = os.path.join(cfg.DATA_PATH,'VOCdevkit', 'VOC2007')
#train_data_path_2012 = os.path.join(cfg.DATA_PATH,'VOCdevkit', 'VOC2012')
train_annotation_path = os.path.join('../data', 'train_annotation.txt')
if os.path.exists(train_annotation_path):
os.remove(train_annotation_path)
# val_set : VOC2007_test
test_data_path_2007 = os.path.join(cfg.DATA_PATH,'VOCdevkit', 'VOC2007')
test_annotation_path = os.path.join('../data', 'test_annotation.txt')
if os.path.exists(test_annotation_path):
os.remove(test_annotation_path)
len_train = parse_voc_annotation(train_data_path_2007, "trainval", train_annotation_path, use_difficult_bbox=False) #+ \
#parse_voc_annotation(train_data_path_2012, "trainval", train_annotation_path, use_difficult_bbox=False)
len_test = parse_voc_annotation(test_data_path_2007, "test", test_annotation_path, use_difficult_bbox=False)
print("The number of images for train and test are :train : {0} | test : {1}".format(len_train, len_test))
内容如下：
路径+x1+y1+x2+y2+类别
D:\research\pytorch-YOLOv4-master\data\VOCdevkit\VOC2007\JPEGImages\000004.bmp 320,287,351,297,0
载入预训练模型
链接：https://pan.baidu.com/s/1NjQ3FKN0gEIKDCDTplw7gw
提取码：4dbe
下载后得到了yolov4.wights文件，这是训练网络所需要的预训练模型
在根目录创建weight，将其放入。
训练
在train.py中，找到如下位置，稍作修改：
if __name__ == "__main__":
global logger, writer
parser = argparse.ArgumentParser()
parser.add_argument('--weight_path', type=str, default='weight/yolov4.weights', help='weight file path')#weight/darknet53_448.weights#修改权重
parser.add_argument('--resume', action='store_true',default=False, help='resume training flag')
parser.add_argument('--gpu_id', type=int, default=-1, help='whither use GPU(eg:0,1,2,3,4,5,6,7,8) or CPU(-1)')
parser.add_argument('--log_path', type=str, default='./log/', help='log path')
parser.add_argument('--accumulate', type=int, default=2, help='batches to accumulate before optimizing')
parser.add_argument('--fp_16', type=bool, default=False, help='whither to use fp16 precision')
训练过程如下：
常见的错误
更新：根据最近浏览到的一些错误问题，重新看了这个版本的代码，并更新作为总结：‘
magic_number = pickle_module.load(f, **pickle_load_args)
_pickle.UnpicklingError: invalid load key, '\x00'.
解决方法：到yolov4_config.py中修改主干网络类型：=，问题的原因在于加载了yolov4的预训练模型，但是选择了Mobilenet-YOLOv4，改为如下即可。
MODEL_TYPE = {
"TYPE": "YOLOv4"#"Mobilenet-YOLOv4"
} # YOLO type:YOLOv4, Mobilenet-YOLOv4 or Mobilenetv3-YOLOv4
2.无法生成模型权重文件：
再train,py中59行左右，验证模型必须在30或50epoch以上才可以验证结果：必须在config设置大于50epoch
self.eval_epoch = (
30 if cfg.MODEL_TYPE["TYPE"] == "YOLOv4" else 50
)
3.not enough values to upack(except 3,got0)
解决方案：很有可能是测试验证数据集没有读入模型，有可能是图片后缀错误，本套代码默认jpg格式，如果其他格式需要自行修改evaluator.py中70附近的后缀。