成功运行SSD_Pytorch并训练自己的数据集

前言

本文主要是记录自己在跑ssd_pytorch算法过程，遇到的一些问题及解决方法。其中大部分的解决方法都是借鉴网上的。写下来主要是方便自己以后查找，也希望能给别人带来点帮助。

一、环境

Windows10

Anaconda3（python3.6.12）

GPU：GTX2060

cuda10.0

pytorch1.7.1

推荐使用pytorch1.3以下的版本，因为源码使用的版本比较低，采用较低版本的pytorch，运行代码可能就不会有那么多的错误。本人采用的版本太高了，所以需改动的地方较多。

二、准备数据集

按照voc2007的格式制作自己的数据集，可以使用labellmg来制作，网上有很多制作数据集的教程；由于我用的是别人制作的数据集，不方便分享。所以在此给出voc2007和voc2012的数据集。

1.voc2007数据集

链接:voc2007

2.voc2012数据集

链接:voc2012

三、从Github上下载作者源码

SSD-Pytorch代码链接:https://github.com/amdegroot/ssd.pytorch
由于SSD的基础网络是vgg16，代码采用的预训练模型，所以在训练时，需先下载vgg16的训练模型。vgg16预训练模型的链接:vgg16

创建文件夹

在根目录下创建weights文件夹，将下载好的预训练vgg16模型文件放在里面。在data目录下创建VOCdevkit文件夹，将下载好的voc2007数据集和voc2012数据集放在里面。训练自己数据集时，将voc2007或voc2012里面的文件换成自己的。

四、修改代码

1.修改data/config.py

主要是修改num_classes(类别数+1)，由于我的数据集只有两类，所以就设为3

# SSD300 CONFIGS
voc = {
   
    'num_classes': 3,  # voc为21，训练自己的数据集时为：类别数+1
    'lr_steps': (80000, 100000, 120000),
    'max_iter': 120000,  # 原120000 最大迭代次,可以自己修改训练的迭代次数
    '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': 'VOC',
}

2.修改data/voc0712.py

修改数据集类别

# 20-25行；注释掉原有的voc数据集的类别，然后根据自己的数据集来修改类别。
VOC_CLASSES = (  # always index 0</