Pytorch 1.0实现SSD

环境 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

可视化结果：

七、修改修改ssd300_voc0712.yaml文件中的DATASETS的train和test名称时一定要注意后面有个逗号，否则会报错。