跑通caffe-ssd demo代码（训练测试自己数据集）
Posted A big bliss
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了跑通caffe-ssd demo代码（训练测试自己数据集）相关的知识，希望对你有一定的参考价值。
跑通caffe-ssd demo代码（训练、测试自己数据集）
ssd网络我就不多介绍了，CSDN上面一搜一大把。这篇主要讲讲如何跑通caffe代码~
github：caffe-ssd
一、代码结构
在caffe-ssd中能用到的文件我全部在上述图片中标出来了，到时候具体的再细说~
caffe-ssd的环境自己百度吧，网上很多安装教程~
二、数据集准备
训练模型，首先第一个事情就是准备数据集。在利用caffe训练分类模型的时候，通常使用lmdb或者hdf5格式的数据，但是在该项目中使用的是lmdb格式的（其他格式的数据肯定也行，但是就是需要自己c++手写数据处理层了~太麻烦了）
1.
首先，准备VOC格式的数据
这是VOC数据集的格式，其中，Annotations里面存放的是所有数据图片相对应的xml标签文件，imagesets里面存放的是一些txt文件，后续用到再详细说，JPEGImages里面存放的就是训练图片，最后两个文件夹是用于实例分割的，在目标检测中用不到，所以就不用管它。（这里，为了简单，所以只使用VOC_trainval_2007中前36张图片）
JPEGImages:
其中，trainval里面存放训练图片，test存放测试图片。这个划分根据你自己项目进行决定就行。一般来说，trainval：test=4:1，这里我在test存放8张随机图片，剩余的存放在trainval文件夹中。值得一提，全部的图片还是放在trainval和test文件夹同级目录中，这样便于后续生成trainval.txt和test.txt。
Annotations:
Imagesets:
2.
下面介绍如何生成Imagesets里面的这四个txt文件~
先说trainval.txt和test.txt这两个文件。
#! /usr/bin/python
# -*- coding:UTF-8 -*-
import os, sys
import glob
trainval_dir = r"D:\\voc\\VOC_test\\JPEGImages\\trainval" #训练集图片存放地址
test_dir = r"D:\\voc\\VOC_test\\JPEGImages\\test" #测试图片存放地址
trainval_img_lists = glob.glob(trainval_dir + '/*.jpg') #如果你的图片是png格式的，只需要修改最后的后缀
trainval_img_names = []
for item in trainval_img_lists:
temp1, temp2 = os.path.splitext(os.path.basename(item))
trainval_img_names.append(temp1)
test_img_lists = glob.glob(test_dir + '/*.jpg') #如果你的图片是png格式的，只需要修改最后的后缀
test_img_names = []
for item in test_img_lists:
temp1, temp2 = os.path.splitext(os.path.basename(item))
test_img_names.append(temp1)
#图片路径和xml路径
dist_img_dir = r"D:\\voc\\VOC_test\\JPEGImages" #JPEGImages路径
dist_anno_dir = r"D:\\voc\\VOC_test\\Annotations" #存放所有数据的xml文件路径
trainval_fd = open(r"D:\\voc\\VOC_test\\ImageSets\\trainval.txt", 'w') #trainval.txt存放地址
test_fd = open(r"D:\\voc\\VOC_test\\ImageSets\\test.txt", 'w') #test.txt存放地址
for item in trainval_img_names:
trainval_fd.write(dist_img_dir + '/' + str(item) + '.jpg' + ' ' + dist_anno_dir + '/' + str(item) + '.xml\\n')
for item in test_img_names:
test_fd.write(dist_img_dir + '/' + str(item) + '.jpg' + ' ' + dist_anno_dir + '/' + str(item) + '.xml\\n')
生成的trainval.txt：
test.txt:
注意：这里我建议大家使用绝对路径，到时候在训练的时候比较清楚点
3.
下面介绍labelmap_voc.prototxt：
labelmap_voc.prototxt文件在你下载的caffe-ssd中有一个副本，位置在：you_caffe_root/data/VOC0712/里面，这里面的数据要根据你自己的需求进行修改，因为我这里就是VOC数据集，所以我不用改变的。不过假如我想进行猫狗目标检测算法，那我就得这么改：
这里，label为0的是背景一类，不管你是检测多少种物体，这一类是不能动的。
4.
下面就是test_name_size.txt生成方式：
#! /usr/bin/python
# -*- coding:UTF-8 -*-
import os, sys
import glob
from PIL import Image #读图
img_dir = r"D:/voc/VOC_test/JPEGImages/test" #测试图片存放路径
#获取制定路径下的所有jpg图片的名称
img_lists = glob.glob(img_dir + '/*.jpg')
test_name_size = open(r"D:/voc/VOC_test/ImageSets/test_name_size.txt", 'w') #test_name_size.txt存放路径
for item in img_lists:
img = Image.open(item)
width, height = img.size
temp1, temp2 = os.path.splitext(os.path.basename(item))
test_name_size.write(temp1 + ' ' + str(height) + ' ' + str(width) + '\\n')
最后生成的test_name_size.txt：
其中，每一列中第一个表示测试图片名称，第二个和第三个表示的是该测试图片的高和宽。
这样，所有准备工作都做完了~用上述所有文件就可以生成lmdb数据了
5.
将VOC文件夹放在you_caffe_root/data中
定位到you_caffe_root/data/VOC0712，下面应该有两个shell脚本文件：create_list.sh和create_data.sh。前者其实就是生成trainval.txt和test.txt，因为我们已经用python脚本生成好了，所以就可以直接用后者来生成lmdb数据了。
create_data.sh:
cur_dir=$(cd $( dirname ${BASH_SOURCE[0]} ) && pwd )
root_dir=$cur_dir/../.. #你安装caffe的根目录
cd $root_dir
redo=1
data_root_dir="" #上述VOC总文件夹的位置
mapfile="" #上述生成的labelmap_voc.prototxt存放位置
anno_type="detection"
db="lmdb"
min_dim=0
max_dim=0
width=0
height=0
extra_cmd="--encode-type=jpg --encoded"
if [ $redo ]
then
extra_cmd="$extra_cmd --redo"
fi
for subset in test trainval
do
python $root_dir/scripts/create_annoset.py #该sh脚本本质上是调用的是scripts/create_annoset.py，你找到create_annoset.py位置就行
--anno-type=$anno_type
--label-map-file=$mapfile
--min-dim=$min_dim
--max-dim=$max_dim
--resize-width=$width
--resize-height=$height
# 很多人最后训练代码出错其实就是最后一行没有设置正确
--check-label
$extra_cmd #这个不用管
$data_root_dir #这个不用管
$root_dir/data/$dataset_name/$subset.txt #这个位置是上述生成的trainval.txt和test.txt文件路径，这个要设置好
$data_root_dir/$dataset_name/$db/$dataset_name"_"$subset"_"$db #这一个参数是设置生成lmdb文件的路径，一般来说，最好设置在VOC总文件夹下面，即同JPEGImages、Annotations、Images在同一个路径下，当然，也可以事先自己在VOC总文件夹下面创建一个lmdb的文件夹，该lmdb文件夹下面又有两个子文件夹，分别代表的是训练lmdb数据和测试lmdb数据
examples/$dataset_name
done
注意点：如果代码运行报错(明明图片和xml标签路径是正确的，但是就是无法生成lmdb)，那么找到create_annoset.py第87行：
img_file, anno = line.strip("\\n").split(" ")
改为
img_file, anno = line.strip("\\n").strip("\\r").split(" ")
三、开始训练
与以往caffe实现分类网络不同的是，该项目是通过py脚本进行训练的，而不是直接通过caffe的c++接口进行训练。
该py脚本的位置在：you_caffe_root/examples/ssd/ssd_pascal.py。
上述图片中还有一个ssd_detect.py脚本，该脚本就是用于测试单张图片用的，这个后续再说。
ssd_pascal.py：
具体需要修改的地方会在下面标注清楚的
from __future__ import print_function
import caffe
from caffe.model_libs import *
from google.protobuf import text_format
import math
import os
import shutil
import stat
import subprocess
import sys
# Add extra layers on top of a "base" network (e.g. VGGNet or Inception).
def AddExtraLayers(net, use_batchnorm=True, lr_mult=1):
use_relu = True
# Add additional convolutional layers.
# 19 x 19
from_layer = net.keys()[-1]
# TODO(weiliu89): Construct the name using the last layer to avoid duplication.
# 10 x 10
out_layer = "conv6_1"
ConvBNLayer(net, from_layer, out_layer, use_batchnorm, use_relu, 256, 1, 0, 1,
lr_mult=lr_mult)
from_layer = out_layer
out_layer = "conv6_2"
ConvBNLayer(net, from_layer, out_layer, use_batchnorm, use_relu, 512, 3, 1, 2,
lr_mult=lr_mult)
# 5 x 5
from_layer = out_layer
out_layer = "conv7_1"
ConvBNLayer(net, from_layer, out_layer, use_batchnorm, use_relu, 128, 1, 0, 1,
lr_mult=lr_mult)
from_layer = out_layer
out_layer = "conv7_2"
ConvBNLayer(net, from_layer, out_layer, use_batchnorm, use_relu, 256, 3, 1, 2,
lr_mult=lr_mult)
# 3 x 3
from_layer = out_layer
out_layer = "conv8_1"
ConvBNLayer(net, from_layer, out_layer, use_batchnorm, use_relu, 128, 1, 0, 1,
lr_mult=lr_mult)
from_layer = out_layer
out_layer = "conv8_2"
ConvBNLayer(net, from_layer, out_layer, use_batchnorm, use_relu, 256, 3, 0, 1,
lr_mult=lr_mult)
# 1 x 1
from_layer = out_layer
out_layer = "conv9_1"
ConvBNLayer(net, from_layer, out_layer, use_batchnorm, use_relu, 128, 1, 0, 1,
lr_mult=lr_mult)
from_layer = out_layer
out_layer = "conv9_2"
ConvBNLayer(net, from_layer, out_layer, use_batchnorm, use_relu, 256, 3, 0, 1,
lr_mult=lr_mult)
return net
### Modify the following parameters accordingly ###
# The directory which contains the caffe code.
# We assume you are running the script at the CAFFE_ROOT.
caffe_root = os.getcwd()
# Set true if you want to start training right after generating all files.
run_soon = True
# Set true if you want to load from most recently saved snapshot.
# Otherwise, we will load from the pretrain_model defined below.
resume_training = True
# If true, Remove old model files.
remove_old_models = False
# The database file for training data. Created by data/VOC0712/create_data.sh
train_data = "examples/VOC0712/VOC0712_trainval_lmdb" #该地址为上述生成的lmdb格式的训练数据，注意的是，该地址为存放lmdb数据的上一级文件夹
# The database file for testing data. Created by data/VOC0712/create_data.sh
test_data = "examples/VOC0712/VOC0712_test_lmdb" #该地址为上述生成的lmdb格式的测试数据，注意的是，该地址为存放lmdb数据的上一级文件夹
# Specify the batch sampler.
resize_width = 300 #该模型接收输入图片的宽
resize_height = 300 #该模型接收输入图片的高
resize = "{}x{}".format(resize_width, resize_height)
batch_sampler = [
{
'sampler': {
},
'max_trials': 1,
'max_sample': 1,
},
{
'sampler': {
'min_scale': 0.3,
'max_scale': 1.0,
'min_aspect_ratio': 0.5,
'max_aspect_ratio': 2.0,
},
'sample_constraint': {
'min_jaccard_overlap': 0.1,
},
'max_trials': 50,
'max_sample': 1,
},
{
'sampler': {
'min_scale': 0.3,
'max_scale': 1.0,
'min_aspect_ratio': 0.5,
'max_aspect_ratio': 2.0,
},
'sample_constraint': {
'min_jaccard_overlap': 0.3,
},
'max_trials': 50,
'max_sample': 1,
},
{
'sampler': {
'min_scale': 0.3,
'max_scale': 1.0,
'min_aspect_ratio': 0.5,
'max_aspect_ratio': 2.0,
},
'sample_constraint': {
'min_jaccard_overlap': 0.5,
},
'max_trials': 50,
'max_sample': 1,
},
{
'sampler': {
'min_scale': 0.3,
'max_scale': 1.0,
'min_aspect_ratio': 0.5,
'max_aspect_ratio': 2.0,
},
'sample_constraint': {
'min_jaccard_overlap': 0.7,
},
'max_trials': 50,
'max_sample': 1,
},
{
'sampler': {
'min_scale': 0.3,
'max_scale': 1.0,
'min_aspect_ratio': 0.5,
'max_aspect_ratio': 2.0,
},
'sample_constraint': {
'min_jaccard_overlap': 0.9,
},
'max_trials': 50,
'max_sample': 1,
},
{
'sampler': {
'min_scale': 0.3,
'max_scale': 1.0,
'min_aspect_ratio': 0.5,
'max_aspect_ratio': 2.0,
},
'sample_constraint': {
'max_jaccard_overlap': 1.0,
},
'max_trials': 50,
'max_sample': 1,
},
]
train_transform_param = { #这个字典里面就是进行数据增强的操作，其实我感觉可以不要该操作，但是没有尝试过注释后代码还能不能跑通
'mirror': True,
'mean_value': [104, 117, 123], #图片均值
'resize_param': {
'prob': 1,
'resize_mode': P.Resize.WARP,
'height': resize_height,
'width': resize_width,
'interp_mode': [
P.Resize.LINEAR,
P.Resize.AREA,
P.Resize.NEAREST,
P.Resize.CUBIC,
P.Resize.LANCZOS4,
],
},
'distort_param': {
'brightness_prob': 0.5,
'brightness_delta': 32,
'contrast_prob': 0.5,
'contrast_lower': 0.5,
'contrast_upper': 1.5,
'hue_prob': 0.5,
'hue_delta': 18,
'saturation_prob': 0.5,
'saturation_lower': 0.5,
'saturation_upper': 1.5,
'random_order_prob': 0.0,
},
'expand_param': {
'prob': 0.5,
'max_expand_ratio': 4.0,
},
'emit_constraint': {
'emit_type': caffe_pb2.EmitConstraint.CENTER,
}
}
test_transform_param = {
'mean_value': [104, 117, 123],
'resize_param': {
'prob': 1,
'resize_mode': P.Resize.WARP,
'height': resize_height,
'width': resize_width,
'interp_mode': [P.Resize.LINEAR],
},
}
# If true, use batch norm for all newly added layers.
# Currently only the non batch norm version has been tested.
use_batchnorm = False
lr_mult = 1
# Use different initial learning rate.
if use_batchnorm:
base_lr = 0.0004
else:
# A learning rate for batch_size = 1, num_gpus = 1.
base_lr = 0.00004 #一般来说，我们会调用这个学习率，但是实际在训练的时候的学习率应该为base_lr * 25，所以说如果想增减学习率，只需要修改此处就可以
# Modify the job name if you want.
job_name = "SSD_{}".format(resize)
# The name of the model. Modify it if you want.
model_name = "VGG_VOC0712_{}".format(job_name)
# Directory which stores the model .prototxt file.
save_dir = "models/VGGNet/VOC0712/{}".format(job_name) #最后生成caffemodel的位置
# Directory which stores the snapshot of models.
snapshot_dir = "models/VGGNet/VOC0712/{}".format(job_name)
# Directory which stores the job script and log file.
job_dir = "jobs/VGGNet/VOC0712/{}".format(job_name)
# Directory which stores the detection results.
output_result_dir = "{}/data/VOCdevkit/results/VOC2007/{}/Main".format(os.environ['HOME'], job_name)
# model definition files.
train_net_file = "{}/train.prototxt".format(save_dir) #train.prototxt位置
test_net_file = "{}/test.prototxt".format(save_dir) #test.prototxt位置
deploy_net_file = "{}/deploy.prototxt".format(save_dir)
solver_file = "{}/solver.prototxt".format(save_dir)
# snapshot prefix.
snapshot_prefix = "{}/{}".format(snapshot_dir, model_name)
# job script path.
job_file = "{}/{}.sh".format(job_dir, model_name)
# Stores the test image names and sizes. Created by data/VOC0712/create_list.sh
name_size_file = "data/VOC0712/test_name_size.txt" #上述生成的test_name_size.txt位置，最后用绝对路径
# The pretrained model. We use the Fully convolutional reduced (atrous) VGGNet.
pretrain_model = "models/VGGNet/VGG_ILSVRC_16_layers_fc_reduced.caffemodel" #预训练模型，会很大程度减少自己项目的训练时间
# Stores LabelMapItem.
label_map_file = "data/VOC0712/labelmap_voc.prototxt" #上述生成的labelmap_voc.prototxt位置，最后用绝对路径
# MultiBoxLoss parameters.
num_classes = 21 #该位置要换成你自己项目的物体类别个数，别忘了要加上背景
share_location = True
background_label_id=0
train_on_diff_gt = True
normalization_以上是关于跑通caffe-ssd demo代码（训练测试自己数据集）的主要内容，如果未能解决你的问题，请参考以下文章
跑通代码---CVPR2020--StegaStamp: Invisible Hyperlinks in Physical Photographs
深度学习之初识篇——小白也能跑通的深度学习万能框架交通标识牌检测