Windows下自己训练数据集，在pytorch上实现基于SSD的物体检测

Windows下自己训练数据集，在pytorch上实现基于SSD的物体检测

（0）环境

Windows10
PyTorch1.4
pycharm

（一）准备数据集

首先新建了VOC2007文件夹，然后再在里面新建三个文件夹，如下图。
其中，JPEGImages用来存放照片，Annotations存放xml文件。
第一步：对图片进行标注。我这里选了五种类别（书包，钟表，手机，水杯，鼠标）各十张图片，总共五十张照片（由于显卡不行，只选了这么多），标注工具为labelImg，教程网上很多，这里不再赘述。标注得到的xml文件存放在Annotations中。

第二步：用make_txt.py文件将数据集分为4个txt文件，用于后续的训练与测试。下图为make_txt.py文件以及分好的四个文件。

make_txt.py文件

import os
import random

trainval_percent = 0.5
train_percent = 0.5
xmlfilepath = './Annotations/'
txtsavepath = './ImageSets/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(txtsavepath + '/trainval.txt', 'w')
ftest = open(txtsavepath + '/test.txt', 'w')
ftrain = open(txtsavepath + '/train.txt', 'w')
fval = open(txtsavepath + '/val.txt', 'w')

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

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

（二）下载源码并修改

第一步：直接去GitHub上下载SSD源码：https://github.com/amdegroot/ssd.pytorch，我是直接dowmload到D盘。解压之后在打开SSD/data文件夹，在这之下新建一个VOCdevkit文件夹，然后把（一）中做好的VOC2007数据集copy到里面。
第二步：对源码进行修改。
用parcharm打开SSD项目。
修改一：data/init.py程序：将第三行的from .coco import COCODetection, COCOAnnotationTransform, COCO_CLASSES, COCO_ROOT, get_label_map 注释掉。因为我们不使用coco数据集，而是使用自己训练的。
修改二：data/config.py程序：第十五行的 num_classes改成自己设定的类别数+1，所以我这里是6。
修改三：data/voc0712.py:第二十行的VOC_CLASSES =改成自己的类别名，我的是bag,clock,phone,cup,mouse；第二十三行的代码改成VOC_ROOT = './data/VOCdevkit/'；第93行改为image_sets=[('2007', 'trainval')]，因为我们只有2007数据集。
修改四：layers/modules/multibox_loss.py程序：在第97行的loss_c[pos] = 0的全面加上一句loss_c = loss_c.view(num, -1)，在第115行N = num_pos.data.sum()的后面加上.double,继续在下面添加
loss_l = loss_l.double()
loss_c = loss_c.double()
这两句代码。
修改五：ssd.py程序：把这里面所有的num_classes的数量都改为类别数+1，我的都是6。
修改六：train.py程序：搜索这里面的data[0]，把这里面的[0]都删掉，只保留data。第84行的这两行#if args.dataset_root == COCO_ROOT: # parser.error('Must specify dataset if specifying dataset_root')注释掉，因为我们不用COCO。
第165行的images, targets = next(batch_iterator)改成

        try:
            images, targets = next(batch_iterator)
        except StopIteration:
            batch_iterator = iter(data_loader)
            images, targets = next(batch_iterator)

第198行的5000我改成500；迭代500次保存一次文件，大家可改可不改。第200行权重的名称自己随意，我改的是weights/ssd300_voc_。

（三）开始训练

开始训练时需要一个预训练文件vgg16_reducedfc.pth，
百度云链接：https://pan.baidu.com/s/1FfqbkSe8spgxzYtCtPBAOQ
提取码：xg4c
大家可自行下载。
下载之后把他放在SSD项目下新建的weights文件夹下，然后就可以进行训练了。
注：训练中途遇到loss=nan的现象，我的解决方法是将train.py中，parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float,中的 default=1e-3改为default=1e-4。
每500次保存一下文件。直到loss降低到1左右时即可。我训练了2000次，在权重weights得到如下所示文件：

我选了最后一次的文件进行测试，测试代码如下：
demo.py文件，放在demo文件夹里，需要自己修改的已标注

import os
import sys
import torch
from torch.autograd import Variable
import numpy as np
import cv2
from ssd import build_ssd
from data import VOC_CLASSES as labels
from matplotlib import pyplot as plt

module_path = os.path.abspath(os.path.join('..'))
if module_path not in sys.path:
    sys.path.append(module_path)

if torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')

net = build_ssd('test', 300, 6)    # 第一处修改：类别+1
# 将预训练的权重加载到数据集上
net.load_weights('../weights/ssd300_voc_2000.pth')  # 第二处修改：使用自己训练好的文件

# 加载多张图像
imgs = '../test_image/'# 第三处修改：改成你自己的文件夹
img_list = os.listdir(imgs)
for img in img_list:
    # 对输入图像进行预处理
    current_img = imgs + img
    image = cv2.imread(current_img)
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    x = cv2.resize(image, (300, 300)).astype(np.float32)
    x -= (104.0, 117.0, 123.0)
    x = x.astype(np.float32)
    x = x[:, :, ::-1].copy()
    x = torch.from_numpy(x).permute(2, 0, 1)

    # 把图片设为变量
    xx = Variable(x.unsqueeze(0))
    if torch.cuda.is_available():
        xx = xx.cuda()
    y = net(xx)

    # 解析 查看结果

    top_k = 10

    plt.figure(figsize=(6, 6))
    colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
    currentAxis = plt.gca()

    detections = y.data
    scale = torch.Tensor(rgb_image.shape[1::-1]).repeat(2)
    for i in range(detections.size(1)):
        j = 0
        while detections[0, i, j, 0] >= 0.6:
            score = detections[0, i, j, 0]
            label_name = labels[i-1]
            display_txt = '%s: %.2f'%(label_name, score)
            print(display_txt)
            pt = (detections[0,i,j,1:]*scale).cpu().numpy()
            coords = (pt[0], pt[1]), pt[2]-pt[0]+1, pt[3]-pt[1]+1
            color = colors[i]
            currentAxis.add_patch(plt.Rectangle(*coords, fill=False, edgecolor=color, linewidth=2))
            currentAxis.text(pt[0], pt[1], display_txt, bbox={'facecolor':color, 'alpha':0.5})
            j += 1
    plt.imshow(rgb_image)
    plt.show()

（四）结果

本文主要参考博主陈纾的方法，十分感谢！