目录
- 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次,所以效果很差。