pytorch_SSD训练评估自己的数据集

1. 模型环境

1. 利用**ssd.pytorch**模型

2. 运行环境 显卡P5000 ubuntu16.04+CUDA9.2
   pytorch1.0.1 , torchvision 0.2.2

2. 数据集处理

2.1 数据格式转换yolo_to_voc

由于开始制作的数据集是为了在yolov3模型下进行训练的yolo格式，所以我们需要进行转换。
原始数据： 640×640 yolo格式
目标数据： 300×300 voc格式

# -*- coding = utf-8 -*-

import copy
from lxml.etree import Element, SubElement, tostring, ElementTree
import cv2
import os

# 为了把yolo格式标签name_value转换成name
# 注意ssd模型中标签name英文字母需要小写
labels_dict = {"0":'a',"1":'b',"2":'c',"3":'d',"4":'e',"5":'f',"6":'g',"7":'h',"8":'i',
               "9":'j',...,"17":'z' }
cnt = 0
txt_path = ''  # yolo存放txt的文件目录
image_path = ''  # 存放图片的文件目录
path = ''  # 存放生成xml的文件目录

xml_head = '''<annotation>
    <folder>VOC2007</folder>
    <filename>{}</filename>
    <source>
        <database>Unknown</database>
    </source>   
    <size>
        <width>{}</width>
        <height>{}</height>
        <depth>{}</depth>
    </size>
    <segmented>0</segmented>
    '''
xml_obj = '''
    <object>        
        <name>{}</name>
        <pose>Rear</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>{}</xmin>
            <ymin>{}</ymin>
            <xmax>{}</xmax>
            <ymax>{}</ymax>
        </bndbox>
    </object>
    '''
xml_end = '''
</annotation>'''


for (root, dirname, files) in os.walk(image_path):  # 遍历图片文件夹
    for ft in files:
        ftxt = ft.replace('jpg', 'txt')  # ft是图片名字+扩展名，将jpg和txt替换
        fxml = ft.replace('jpg', 'xml')
        xml_path = path + fxml
        obj = ''

        img = cv2.imread(root + ft)
        img_h, img_w = img.shape[0], img.shape[1]
        head = xml_head.format(str(ft), str(img_w), str(img_h), 3)

        with open(txt_path + ftxt, 'r') as f:  # 读取对应txt文件内容
            for line in f.readlines():
                yolo_datas = line.strip().split(' ')
                label = labels_dict[str(yolo_datas[0].strip())]
                center_x = round(float(str(yolo_datas[1]).strip()) * img_w)
                center_y = round(float(str(yolo_datas[2]).strip()) * img_h)
                bbox_width = round(float(str(yolo_datas[3]).strip()) * img_w)
                bbox_height = round(float(str(yolo_datas[4]).strip()) * img_h)

                xmin = str(int(center_x - bbox_width / 2))
                ymin = str(int(center_y - bbox_height / 2))
                xmax = str(int(center_x + bbox_width / 2))
                ymax = str(int(center_y + bbox_height / 2))

                obj += xml_obj.format(label ,xmin, ymin, xmax, ymax)
        with open(xml_path, 'w') as f_xml:
            f_xml.write(head + obj + xml_end)
        cnt += 1
        print(cnt)

2.2 数据划分

按照 train:val:test = 7:2:1

# -*- coding = utf-8 -*-

import os
import random

trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = '/home/geosis03/airplane/images/Annotations/'
txtsavepath = '/home/geosis03//airplane/images/Main/'
total_xml = os.listdir(xmlfilepath)

num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftrainval = open('./VOC2007/ImageSets/Main/trainval.txt', 'w')
ftest = open('./VOC2007/ImageSets/Main/test.txt', 'w')
ftrain = open('./VOC2007/ImageSets/Main/train.txt', 'w')
fval = open('./VOC2007/ImageSets/Main/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

2.3 数据集结构

1. 根据voc2007的数据来构建自己的数据集结构

• VOCdevkit
  • VOC2007
    • Annotations
    • ImageSets
      + Main
    • JPEGImages

Main文件夹下包含：train.txt val.txt test.txt trainval.txt testval.txt
可以根据具体数据情况来设置训练集与测试集情况
我们这里 训练集train.txt 测试集testval.txt 按7：3来划分，与原模型中默认设置不同

3. 模型代码配置

3.1 代码配置

设置思路
SSD-PyTorch训练自己的数据集问题汇总

3.2 设置自己的代码

根据上面的大家设置的思路基本上可以把代码环境配置完毕，然后进一步根据自己的数据集具体情况进行微调。

1. 保留原始的voc0712.py ，创建自己的相应数据集的airplane.py
   修改类别信息等需要把相应的VOC换成AIRPLANE

AIRPLANE_CLASSES = (  # always index 0
    'a', 'b', 'c', ..., 'z')

2.由于我是单独创建了airplane.py来代替voc0712.py，所以需要修改一系列的数据目录，以及import的信息
在 test.py,eval.py,ssd.py中修改 import信息

config.py修改

HOME = os.path.expanduser("/home/geosis03/XXX/model/ssd.pytorch/")

airplane = {
    'num_classes': 19,
    'lr_steps': (28800,32400, 36000),
    'max_iter': 36000,
    'feature_maps': [38, 19, 10, 5, 3, 1],
    'min_dim': 300,
    'steps': [8, 16, 32, 64, 100, 300],
    'min_sizes': [30, 60, 111, 162, 213, 264],
    'max_sizes': [60, 111, 162, 213, 264, 315],
    'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
    'variance': [0.1, 0.2],
    'clip': True,
    'name': 'AIRPLANE',
}

train.py
需要提前下载与训练的VGG模型
vgg16_reducedfc.pth
train.py中的训练参数设置，详细设置请根据上面的链接进行修改。

parser = argparse.ArgumentParser(
    description='Single Shot MultiBox Detector Training With Pytorch')
train_set = parser.add_mutually_exclusive_group()
parser.add_argument('--dataset', default='AIRPLANE', choices=['VOC', 'COCO','AIRPLANE'],
                    type=str, help='VOC or COCO or AIRPLANE')
parser.add_argument('--dataset_root', default=AIRPLANE_ROOT,
                    help='Dataset root directory path')
parser.add_argument('--basenet', default='vgg16_reducedfc.pth',
                    help='Pretrained base model')
parser.add_argument('--batch_size', default=32, type=int,
                    help='Batch size for training')
parser.add_argument('--resume', default=None, type=str,
                    help='Checkpoint state_dict file to resume training from')
parser.add_argument('--start_iter', default=0, type=int,
                    help='Resume training at this iter')
parser.add_argument('--num_workers', default=4, type=int,
                    help='Number of workers used in dataloading')
parser.add_argument('--cuda', default=True, type=str2bool,
                    help='Use CUDA to train model')
parser.add_argument('--lr', '--learning-rate', default=5e-4, type=float,
                    help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float,
                    help='Momentum value for optim')
parser.add_argument('--weight_decay', default=5e-6, type=float,
                    help='Weight decay for SGD')
parser.add_argument('--gamma', default=0.1, type=float,
                    help='Gamma update for SGD')
parser.add_argument('--visdom', default=False, type=str2bool,
                    help='Use visdom for loss visualization')
parser.add_argument('--save_folder', default='weights/36000/',
                    help='Directory for saving checkpoint models')
args = parser.parse_args()

4. BUG问题

1. 若数据标签名称出现大写字母,引起 key error
   解决办法：
   修改eval.py 286行代码

#将含有小字母的标签名称转换成大写再与标签label中的classname比较
# R = [obj for obj in recs[imagename] if obj['name'] == classname]
for imagename in imagenames:
    if classname != 'aerobus':
        R = [obj for obj in recs[imagename] if obj['name'] == classname.upper()]
    else:
        R = [obj for obj in recs[imagename] if obj['name'] == classname]

2. 若要修改训练数据集，将默认的trainval.txt换成train.txt
   则在airplane.py中修改

    def __init__(self, root,
                 image_sets='train',
                 transform=None, target_transform=AIRPLANEAnnotationTransform(),
                 dataset_name='AIRPLANE'):

若要修改测试集将test.txt 改为testval.txt
修改test.py 87行

testset = AIRPLANEDetection(args.voc_root, 'testval', None, AIRPLANEAnnotationTransform())

修改eval.py 75行

set_type = 'testval'

最后还需要将数据集
VOC2007文件夹中annots.pkl删除，这样完成了更改测试集的步骤

3. 如要可视化显示
   首先安装visdom
   然后需要修改train.py 95行附近

    if args.visdom:
        import visdom
        global viz
        viz = visdom.Visdom()

最后，运行该模型代码可能还有些许小问题，需要耐心按步骤修改。
利用网络上的一些建议方案可能还是不能顺利实现的时候，你就需要仔细考虑一下是否是自身的数据问题。