Caffe 中LMDB接口处理多标签（Multi-Label）数据

布满荆棘的人生 2022-08-20 10:30 151阅读 0赞

Caffe的数据接口主要有原始图像（ImageData）, HDF5, LMDB/LevelDB。由于Caffe自带的图像转LMDB接口只支持但label，对于多label的任务，往往需要使用HDF5。

然而，Caffe对于HDF5的数据，需要预先读取整个h5文件，这对于小数据的情况不成问题，而且一次性读到内存里还节省训练中的IO开销。可是对于数据量大的情况，内存可能放不下整个h5文件，则需要划分成几个小的h5文件。可这样的实现一方面不优雅，另一方面训练中需要不停地轮流读取h5文件。一种可能的解决方案是把图像数据放到lmdb，label数据放到h5文件，prototxt里面label和data分别来自两个data layer。可是个人觉得这样的实现也不好看，毕竟代码里面要做HDF5和LMDB的存储。

最近从[网上][Link 1]看到一种更直接的方法，大致是结合Python的LMDB库和Caffe的Python 接口 caffe.io.array\_to\_datum，把图像数据和label,分别存储到两个lmdb文件。而对于存储好的lmdb，又怎样写prototxt里面的datalayer来读取呢？目前caffe的datalayer, 指明了LMDB作为backend的话，默认第一个top就是存储lmdb时datum的data，第二个top就是datum的label，在下面的代码里没有指定datum的label，因此，对于data和label的lmdb，分别写一个datalayer, 每个datalayer的第一个top就是对应lmdb里的内容了。而top的blob的名字是可以自己定义的。  
代码如下：

def write_lmdb(image_name_list,label_array,lmdb_img_name,lmdb_label_name,resize_image = False):
        for lmdb_name in [lmdb_img_name, lmdb_label_name]:
            db_path = os.path.abspath(lmdb_name)
            if os.path.exists(db_path):
                shutil.rmtree(db_path)
        counter_img = 0
        counter_label = 0
        batchsz = 100
        fail_cnt = 0
        print("Processing {:d} images and labels...".format(len(image_name_list)))
        for i in xrange(int(np.ceil(len(image_name_list)/float(batchsz)))):
            image_name_batch = image_name_list[batchsz*i:batchsz*(i+1)]
            label_batch = label_array[batchsz*i:batchsz*(i+1),:]
            print label_batch[np.newaxis,np.newaxis,0].dtype
            raw_input('r')
            imgs, labels = [], []
            for idx,image_name in enumerate(image_name_batch):
                    img = skimage.io.imread(image_name)
                    if resize_image==True:
                        img = skimage.transform.resize(img,(96,96))
                    imgs.append(img)
    
            db_imgs = lmdb.open(lmdb_img_name, map_size=1e12)
            with db_imgs.begin(write=True) as txn_img:
                for img in imgs:
                    datum = caffe.io.array_to_datum(np.expand_dims(img, axis=0))
                    txn_img.put("{:0>10d}".format(counter_img), datum.SerializeToString())
                    counter_img += 1
            print("Processed {:d} images".format(counter_img))
            db_labels = lmdb.open(lmdb_label_name, map_size=1e12)
            with db_labels.begin(write=True) as txn_label:
                for idx in range(label_batch.shape[0]):
                    datum = caffe.io.array_to_datum(label_batch[np.newaxis,np.newaxis,idx])
                    txn_label.put("{:0>10d}".format(counter_label), datum.SerializeToString())
                    counter_label += 1
            print("Processed {:d} labels".format(counter_label))
        print fail_cnt,'images fail reading'
        db_imgs.close()
        db_labels.close()

[Link 1]: https://gist.github.com/ajsander/b65061d12f50de3cef5d#file-fcn_tutorial-ipynb