SSD-Tensorflow:基于TensorFlow实现的单发多框检测器

SSD(Single Shot MultiBox Detector)是一种统一的目标检测框架,使用单个神经网络就可以完成目标检测任务。它最初由Wei Liu等人在2016年的一篇研究论文中提出。

SSD-Tensorflow项目是对原始Caffe实现的TensorFlow重新实现。目前,该项目主要实现了基于VGG的SSD网络(输入尺寸为300和512),但项目的架构是模块化的,这使得实现和训练其他SSD变体(如基于ResNet或Inception的网络)变得容易。当前的TensorFlow检查点是直接从SSD Caffe模型转换而来的。

项目结构

该项目的组织结构受到TF-Slim模型仓库的启发,后者包含了许多流行架构(ResNet、Inception和VGG)的实现。SSD-Tensorflow主要分为三个部分:

1. datasets:用于处理常用数据集(如Pascal VOC、COCO等)的接口,以及将数据集转换为TF-Records格式的脚本。

2. networks:SSD网络的定义,以及常用的编码和解码方法。

3. pre-processing:预处理和数据增强例程,参考了原始VGG和Inception实现。

SSD最小示例

SSD Notebook包含了SSD TensorFlow流程的最小示例。简而言之,检测过程包括两个主要步骤:

1. 在图像上运行SSD网络
2. 使用常见算法(如top-k过滤和非最大抑制算法)对输出进行后处理

以下是两个成功检测输出的示例:

要运行该notebook,首先需要解压checkpoint文件:

unzip ssd_300_vgg.ckpt.zip

然后启动Jupyter notebook:

jupyter notebook notebooks/ssd_notebook.ipynb

数据集

当前版本仅支持Pascal VOC数据集(2007和2012年)。为了用于训练SSD模型,需要使用tf_convert_data.py脚本将数据集转换为TF-Records格式:

DATASET_DIR=./VOC2007/test/
OUTPUT_DIR=./tfrecords
python tf_convert_data.py \
    --dataset_name=pascalvoc \
    --dataset_dir=${DATASET_DIR} \
    --output_name=voc_2007_train \
    --output_dir=${OUTPUT_DIR}

注意,上述命令会生成一系列TF-Records文件,而不是单个文件,这是为了在训练过程中更容易进行数据洗牌。

Pascal VOC 2007评估

当前TensorFlow实现的SSD模型性能如下:

我们正在努力复现原始Caffe实现的性能!

下载并解压上述检查点后,可以通过运行以下命令来重现评估指标:

EVAL_DIR=./logs/
CHECKPOINT_PATH=./checkpoints/VGG_VOC0712_SSD_300x300_ft_iter_120000.ckpt
python eval_ssd_network.py \
    --eval_dir=${EVAL_DIR} \
    --dataset_dir=${DATASET_DIR} \
    --dataset_name=pascalvoc_2007 \
    --dataset_split_name=test \
    --model_name=ssd_300_vgg \
    --checkpoint_path=${CHECKPOINT_PATH} \
    --batch_size=1

评估脚本会提供召回率-精确率曲线的估计,并按照Pascal VOC 2007和2012的指南计算mAP指标。

此外,如果想要测试不同的Caffe SSD检查点,可以使用以下命令将其转换为TensorFlow检查点:

CAFFE_MODEL=./ckpts/SSD_300x300_ft_VOC0712/VGG_VOC0712_SSD_300x300_ft_iter_120000.caffemodel
python caffe_to_tensorflow.py \
    --model_name=ssd_300_vgg \
    --num_classes=21 \
    --caffemodel_path=${CAFFE_MODEL}

训练

train_ssd_network.py脚本负责训练网络。与TF-Slim模型类似,可以向训练过程传递多个选项(如数据集、优化器、超参数、模型等)。特别地,可以提供一个检查点文件作为起点,以便微调网络。

微调现有SSD检查点

微调SSD模型最简单的方法是使用预训练的SSD网络(VGG-300或VGG-512)。例如,可以从以下命令开始微调模型:

DATASET_DIR=./tfrecords
TRAIN_DIR=./logs/
CHECKPOINT_PATH=./checkpoints/ssd_300_vgg.ckpt
python train_ssd_network.py \
    --train_dir=${TRAIN_DIR} \
    --dataset_dir=${DATASET_DIR} \
    --dataset_name=pascalvoc_2012 \
    --dataset_split_name=train \
    --model_name=ssd_300_vgg \
    --checkpoint_path=${CHECKPOINT_PATH} \
    --save_summaries_secs=60 \
    --save_interval_secs=600 \
    --weight_decay=0.0005 \
    --optimizer=adam \
    --learning_rate=0.001 \
    --batch_size=32

除了训练脚本标志外,还可以在ssd_vgg_preprocessing.py中尝试数据增强参数(如随机裁剪、分辨率等),或在ssd_vgg_300/512.py中调整网络参数(如特征层、锚框等)。

此外,训练脚本可以与评估例程结合使用,以监控保存的检查点在验证数据集上的性能。为此,可以为训练和验证脚本设置GPU内存上限,使它们能够在同一设备上并行运行。如果有可用的GPU内存用于评估脚本,可以按如下方式并行运行:

EVAL_DIR=${TRAIN_DIR}/eval
python eval_ssd_network.py \
    --eval_dir=${EVAL_DIR} \
    --dataset_dir=${DATASET_DIR} \
    --dataset_name=pascalvoc_2007 \
    --dataset_split_name=test \
    --model_name=ssd_300_vgg \
    --checkpoint_path=${TRAIN_DIR} \
    --wait_for_checkpoints=True \
    --batch_size=1 \
    --max_num_batches=500

微调在ImageNet上预训练的网络

也可以尝试基于标准架构(如VGG、ResNet、Inception等)构建新的SSD模型,并在其上设置multibox层(具有特定的锚框、比例等)。为此,可以通过仅加载原始架构的权重来微调网络,并随机初始化网络的其余部分。例如,对于VGG-16架构,可以按如下方式训练新模型:

DATASET_DIR=./tfrecords
TRAIN_DIR=./log/
CHECKPOINT_PATH=./checkpoints/vgg_16.ckpt
python train_ssd_network.py \
    --train_dir=${TRAIN_DIR} \
    --dataset_dir=${DATASET_DIR} \
    --dataset_name=pascalvoc_2007 \
    --dataset_split_name=train \
    --model_name=ssd_300_vgg \
    --checkpoint_path=${CHECKPOINT_PATH} \
    --checkpoint_model_scope=vgg_16 \
    --checkpoint_exclude_scopes=ssd_300_vgg/conv6,ssd_300_vgg/conv7,ssd_300_vgg/block8,ssd_300_vgg/block9,ssd_300_vgg/block10,ssd_300_vgg/block11,ssd_300_vgg/block4_box,ssd_300_vgg/block7_box,ssd_300_vgg/block8_box,ssd_300_vgg/block9_box,ssd_300_vgg/block10_box,ssd_300_vgg/block11_box \
    --trainable_scopes=ssd_300_vgg/conv6,ssd_300_vgg/conv7,ssd_300_vgg/block8,ssd_300_vgg/block9,ssd_300_vgg/block10,ssd_300_vgg/block11,ssd_300_vgg/block4_box,ssd_300_vgg/block7_box,ssd_300_vgg/block8_box,ssd_300_vgg/block9_box,ssd_300_vgg/block10_box,ssd_300_vgg/block11_box \
    --save_summaries_secs=60 \
    --save_interval_secs=600 \
    --weight_decay=0.0005 \
    --optimizer=adam \
    --learning_rate=0.001 \
    --batch_size=32

在上述命令中,训练脚本会随机初始化属于checkpoint_exclude_scopes的权重,并从检查点文件vgg_16.ckpt加载网络的其余部分。注意,我们还使用trainable_scopes参数指定首先只训练新的SSD组件,保持VGG网络的其余部分不变。一旦网络收敛到一个不错的初始结果(例如,约0.5 mAP),就可以按如下方式微调整个网络:

DATASET_DIR=./tfrecords
TRAIN_DIR=./log_finetune/
CHECKPOINT_PATH=./log/model.ckpt-N
python train_ssd_network.py \
    --train_dir=${TRAIN_DIR} \
    --dataset_dir=${DATASET_DIR} \
    --dataset_name=pascalvoc_2007 \
    --dataset_split_name=train \
    --model_name=ssd_300_vgg \
    --checkpoint_path=${CHECKPOINT_PATH} \
    --checkpoint_model_scope=vgg_16 \
    --save_summaries_secs=60 \
    --save_interval_secs=600 \
    --weight_decay=0.0005 \
    --optimizer=adam \
    --learning_rate=0.00001 \
    --learning_rate_decay_factor=0.94 \
    --batch_size=32

许多流行深度架构的预训练权重可以在TF-Slim模型页面上找到。

SSD-Tensorflow项目为研究人员和开发者提供了一个灵活的框架,用于在TensorFlow中实现和训练SSD目标检测模型。通过提供详细的说明和示例,该项目使得在各种数据集上训练和评估SSD模型变得简单。无论是微调现有模型还是从头开始训练新模型,SSD-Tensorflow都提供了必要的工具和脚本来支持这些工作流程。