环境 ubantu16.04+cudnn7.0+cuda_9.0.176
Pytorch1.0+python3.6.5+ anaconda3
源代码:(https://github.com/lufficc/SSD)
一、数据准备
1、VOC数据集
2、下载预训练权重, vgg16_reducedfc.pth
二、修改文件
1、修改ssd300_voc0712.yaml文件
修改NUM_CLASSES为要训练的类别。数据集train和test的名称。MAX_ITER, LR_STEPS,等参数
MODEL:
NUM_CLASSES: 3
INPUT:
IMAGE_SIZE: 300
DATASETS:
TRAIN: ("voc_2007_trainval",)
TEST: ("voc_2007_test", )
SOLVER:
MAX_ITER: 12000
LR_STEPS: [8000, 10000]
GAMMA: 0.1
BATCH_SIZE: 32
LR: 1e-3
2、其中path_catlog.py中的数据集路径如下,按要求存放数据集或按照数据集修改存放的路径
class DatasetCatalog:
DATA_DIR = 'datasets'
DATASETS = {
'voc_2007_train': {
"data_dir": "VOC2007",
"split": "train"
},
'voc_2007_val': {
"data_dir": "VOC2007",
"split": "val"
},
'voc_2007_trainval': {
"data_dir": "VOC2007",
"split": "trainval"
},
'voc_2007_test': {
"data_dir": "VOC2007",
"split": "test"
},
}
3、修改default.py文件
修改NUM_CLASSES,_C.SOLVER.MAX_ITER ,_C.SOLVER.LR_STEPS 等参数
from yacs.config import CfgNode as CN
_C = CN()
_C.MODEL = CN()
_C.MODEL.DEVICE = "cuda"
# match default boxes to any ground truth with jaccard overlap higher than a threshold (0.5)
_C.MODEL.THRESHOLD = 0.5
_C.MODEL.NUM_CLASSES = 3 #21
# Hard negative mining..+
_C.MODEL.NEG_POS_RATIO = 3
_C.MODEL.CENTER_VARIANCE = 0.1
_C.MODEL.SIZE_VARIANCE = 0.2
# -----------------------------------------------------------------------------
# PRIORS
# -----------------------------------------------------------------------------
_C.MODEL.PRIORS = CN()
_C.MODEL.PRIORS.FEATURE_MAPS = [38, 19, 10, 5, 3, 1]
_C.MODEL.PRIORS.STRIDES = [8, 16, 32, 64, 100, 300]
_C.MODEL.PRIORS.MIN_SIZES = [30, 60, 111, 162, 213, 264]
_C.MODEL.PRIORS.MAX_SIZES = [60, 111, 162, 213, 264, 315]
_C.MODEL.PRIORS.ASPECT_RATIOS = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
# When has 1 aspect ratio, every location has 4 boxes, 2 ratio 6 boxes.
# #boxes = 2 + #ratio * 2
_C.MODEL.PRIORS.BOXES_PER_LOCATION = [4, 6, 6, 6, 4, 4] # number of boxes per feature map location
_C.MODEL.PRIORS.CLIP = True
# -----------------------------------------------------------------------------
# INPUT
# -----------------------------------------------------------------------------
_C.INPUT = CN()
# Image size
_C.INPUT.IMAGE_SIZE = 300
# Values to be used for image normalization, RGB layout
_C.INPUT.PIXEL_MEAN = [123, 117, 104]
# -----------------------------------------------------------------------------
# Dataset
# -----------------------------------------------------------------------------
_C.DATASETS = CN()
# List of the dataset names for training, as present in paths_catalog.py
_C.DATASETS.TRAIN = ()
# List of the dataset names for testing, as present in paths_catalog.py
_C.DATASETS.TEST = ()
# ---------------------------------------------------------------------------- #
# Solver
# ---------------------------------------------------------------------------- #
_C.SOLVER = CN()
# train configs
_C.SOLVER.MAX_ITER = 12000
_C.SOLVER.LR_STEPS = [8000, 10000]
_C.SOLVER.GAMMA = 0.1
_C.SOLVER.BATCH_SIZE = 32
_C.SOLVER.LR = 1e-3
_C.SOLVER.MOMENTUM = 0.9
_C.SOLVER.WEIGHT_DECAY = 5e-4
_C.SOLVER.WARMUP_FACTOR = 1.0 / 3
_C.SOLVER.WARMUP_ITERS = 500
# ---------------------------------------------------------------------------- #
# Specific test options
# ---------------------------------------------------------------------------- #
_C.TEST = CN()
_C.TEST.NMS_THRESHOLD = 0.45
_C.TEST.CONFIDENCE_THRESHOLD = 0.01
# change MAX_PER_CLASS to 400 as official caffe code will slightly increase mAP(0.8025=>0.8063, 0.7783=>0.7798)
_C.TEST.MAX_PER_CLASS = 200
_C.TEST.MAX_PER_IMAGE = -1
_C.OUTPUT_DIR = 'output'
三、训练
输入命令:
$ python train_ssd.py --config-file configs/ssd300_voc0712.yaml --vgg vgg16_reducedfc.pth
训练结果:
训练过后自动测试验证集,并在output文件夹中生成文件,如下所示:
四、验证:
$ python eval_ssd.py --config-file configs/ssd300_voc0712.yaml --weights output/
ssd300_vgg_final.pth
五、测试:
运行命令:
$ python demo.py --config-file configs/ssd300_voc0712.yaml --weights output/
ssd300_vgg_final.pth --images_dir demo
运行结果:
六、可视化
运行命令:
$ tensorboard --logdir=output
可视化结果: