ssd.pytorch指南（五）——训练自己的数据集

目录

• ssd.pytorch指南（一）——序言
• ssd.pytorch指南（二）——平台搭建
• ssd.pytorch指南（三）——训练模型
• ssd.pytorch指南（四）——制作PASCAL VOC2007数据集
• ssd.pytorch指南（五）——训练自己的数据集

在完成制作PASCAL VOC2007数据集后可以开始训练自己的数据集了。

一、修改参数
1、 修改config.py参数
　　打开ssd_pytorch/data/config.py，可以找到这样一段代码，要修改的就是num_class(类别数）+1，根据自己数据集的实际情况进行修改，因为除了原有的分类数，还会自动将背景划为1类。例如，原始数据库的有20中分类，因此’num_classes’:为21。我的数据库为2类，则’num_classes’:为3。一定得是类别数+1，不然会报错！！！。
　　至于max_iter指的是最大迭代数，在尝试训练阶段，建议把它调低，比如调到20，在确认训练过程没有问题之后，在调大迭代次数。这样可以有效的节约训练时间。

voc = {
    'num_classes': 21,
    'lr_steps': (80000, 100000, 120000),
    'max_iter': 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、 修改voc0712.py参数
　　因为我是用的VOC格式的数据集，所以要在ssd_pytorch/data/voc0712.py处理数据。按照自己的数据的类别进行修改。

VOC_CLASSES = (  # always index 0
    'aeroplane', 'bicycle', 'bird', 'boat',
    'bottle', 'bus', 'car', 'cat', 'chair',
    'cow', 'diningtable', 'dog', 'horse',
    'motorbike', 'person', 'pottedplant',
    'sheep', 'sofa', 'train', 'tvmonitor')

3、 修改ssd.py参数
　　这一步还是修改代码中的类别数。找到ssd_pytorch/ssd.py按照自己的数据的num_classes类别数进行修改。大概第32行。

    def __init__(self, phase, size, base, extras, head, num_classes):
        super(SSD, self).__init__()
        self.phase = phase
        self.num_classes = num_classes
        self.cfg = (coco, voc)[num_classes == 21]           #这里
        self.priorbox = PriorBox(self.cfg)
        self.priors = Variable(self.priorbox.forward(), volatile=True)
        self.size = size

还有第198行

def build_ssd(phase, size=300, num_classes=21):          #这里
    if phase != "test" and phase != "train":
        print("ERROR: Phase: " + phase + " not recognized")
        return
    if size != 300:
        print("ERROR: You specified size " + repr(size) + ". However, " +
              "currently only SSD300 (size=300) is supported!")
        return
    base_, extras_, head_ = multibox(vgg(base[str(size)], 3),
                                     add_extras(extras[str(size)], 1024),
                                     mbox[str(size)], num_classes)
    return SSD(phase, size, base_, extras_, head_, num_classes)

（注：实际类别数+1）
二、训练模型

python train.py

三、测试效果
1、直接运行代码。

python eval.py

2、因为我只训练了20次，所以效果很差。