Caffe中LMDB接口实现多标签数据准备及训练

最新推荐文章于 2024-08-18 12:39:23 发布

hyman_yx

最新推荐文章于 2024-08-18 12:39:23 发布

阅读量1.1w

点赞数 6

分类专栏： Deep Learning 文章标签： Caffe LMDB 多标签 multilable

本文链接：https://blog.csdn.net/hyman_yx/article/details/51791136

版权

Deep Learning 专栏收录该内容

7 篇文章 0 订阅

订阅专栏

有不少博客讲Caffe多标签输入的问题，但总觉得讲的不够透彻，在实践角度上没有给出详细的指导，所以本文力求能给出详细的实践过程和说明。

Caffe多标签输入常用的的方法有以下几种：
1. 修改Caffe源码使其支持多标签输入，参考CSDN博客《caffe 实现多标签输入（multilabel、multitask）》
2. HDF5 + Slice Layer，HDF5支持多标签，但Caffe在读取HDF5时会将所有数据一次性预读进内存中，在数据量较大时就对内存有较高要求了；也可以对数据分片，通过在prototxt文件中加入Slice Layer层使Caffe依次读取每个数据分片。
3. 使用两个data输入（例如两个LMDB，一个存储图片，一个存储多标签），然后修改prototxt文件配置两个data layer。

本文只讨论上面第三种方法支持多标签输入。Caffe中单标签LMDB的创建可以通过自带的脚本很方便的实现，但多标签LMDB的创建并不像单标签那样简单易用，需要使用者对Caffe的Data Layer有足够的了解，对于大多数将Caffe作为一个黑盒的使用者而言，这个问题是非常麻烦的。不过该问题已经纳入了Caffe developer的讨论中了，相信不久后Caffe就能方便快速的进行多标签数据的准备了。

多标签的场景下，数据是 N x H x W x C 的一个4维的blob，对应的标签是 N x M x 1 x 1的一个4维的blob。我们通常通过如下几个步骤来进行数据准备：

创建图片LMDB

使用下面的代码创建两个图片LMDB：
- train_data_lmdb - 用于训练的图片LMDB
- val_data_lmdb - 用于测试的图片LMDB

# 生成训练图片列表文件，即将tarin_images_dir目录下所有图片的路径写入temp.txt文件
find tarin_images_dir -type f -exec echo {} \; > temp.txt

# 在temp.txt文件中每一行后追加伪标签，伪标签起占位符作用，实际并不使用 
sed "s/$/ 0/" temp.txt > train_images.txt

# 根据train_images.txt创建train_data_lmdb
$CAFFE_HOME/build/tools/convert_imageset -resize_height=256 -resize_width=256 / train_images.txt train_data_lmdb

# 计算图片均值
$CAFFE_HOME/build/tools/compute_image_mean train_data_lmdb mean.binaryproto

上面的代码用来生成train_data_lmdb，接下来修改代码生成val_data_lmdb.

创建标签LMDB

使用下面的代码创建两个标签LMDB：
- train_label_lmdb - 用于训练的标签LMDB
- val_label_lmdb - 用于测试的标签LMDB

读入每幅图片的多标签，生成标签LMDB，注意图片LMDB和标签LMDB中的顺序一致，创建LMDB的部分python
代码如下：

import sys 
import numpy as np
import lmdb

caffe_root = '/home/xinxin/caffe/' 
sys.path.insert(0, caffe_root + 'python')
import caffe

# 根据多标签的位置选择从数据库、文件等中读取每幅图片的多标签，将其构造成一维的np.array类型，并追加入all_labels列表
all_labels = []
# Add your code of reading labels here ！

# 创建标签LMDB
key = 0
lmdb_path = "/home/xinxin/dl-vsearch/2016-6-27/train_label_lmdb"
env = lmdb.open(lmdb_path, map_size=map_size)
with env.begin(write=True) as txn: 
    for labels in all_labels:
        datum = caffe.proto.caffe_pb2.Datum()
        datum.channels = labels.shape[0]
        datum.height = 1
        datum.width =  1
        datum.data = labels.tostring()          # or .tobytes() if numpy < 1.9 
        datum.label = 0
        key_str = '{:08}'.format(key)

        txn.put(key_str.encode('ascii'), datum.SerializeToString())
        key += 1

上面的代码用来生成train_label_lmdb，接下来修改代码生成val_label_lmdb.

修改prototxt文件

上面我们已经创建了四个数据集train_data_lmdb、val_data_lmdb、train_label_lmdb、val_label_lmdb，并且计算了图片均值mean.binaryproto。

接下来就是修改神经网络deploy.prototxt使其支持多标签训练。由于有四个数据集，我们只需要在deploy.txt中加入四个data layers，分别用来配置这四个数据集。需要注意的是在标签lmdb对应的data layer的参数transform_param中通过scaling参数对label进行缩放，该操作的作用是按照Sigmoid Cross Entropy Loss函数的要求将label范围从[0,255]正则化到[0,1]（根据采用的多标签Loss函数的要求，scaling参数可以修改)。

将deploy.prototxt中原有的data layers修改为如下的四个data layers：

# ------------- 配置训练阶段的图片数据集 ----------------
layers {
  name: "data"
  type: DATA
  top: "data"
  include {
    phase: TRAIN
  }
  transform_param {
    mirror: true
    crop_size: 224
    mean_file: "data.binaryproto"    # 修改为你的均值文件路径
  }
  data_param {
    source: "train_data_lmdb"        # 修改为你的图片lmdb路径
    batch_size: 32
    backend: LMDB
  }
}

# ------------- 配置训练阶段的标签数据集 ----------------
layers {
  name: "data"
  type: DATA
  top: "label"
  include {
    phase: TRAIN
  }
  transform_param {
    scale: 0.00390625                 # 根据需求设置标签缩放系数
    mean_value: 0
  }
  data_param {
    source: "train_label_lmdb"        # 修改为训练集图片lmdb路径
    batch_size: 32
    backend: LMDB
  }
}

# ------------- 配置测试阶段的图片数据集 ----------------
layers {
  name: "data"
  type: DATA
  top: "data"
  include {
    phase: TEST
  }
  transform_param {
    mirror: false
    crop_size: 224
    mean_file: "mean.binaryproto"      # 修改为你的均值文件路径
  }
  data_param {
    source: "val_data_lmdb"            # 修改为训练阶段的图片lmdb路径
    batch_size: 1
    backend: LMDB
  }
}

# ------------- 配置测试阶段的标签数据集 ----------------
layers {
  name: "data"
  type: DATA
  top: "label"
  include {
    phase: TEST
  }
  transform_param {
    scale: 0.00390625                  # 根据需求设置标签缩放系数
    mean_value: 0
  }
  data_param {
    source: "val_label_lmdb"           # 修改为测试阶段的标签lmdb路径
    batch_size: 1
    backend: LMDB
  }
}

......

# ----------------- 多标签损失函数  -------------------
layers {
  name: "loss"
  type: SIGMOID_CROSS_ENTROPY_LOSS
  bottom: "fc8"                       # 根据需求配置
  bottom: "label"
  top: "loss"
}