将图片文件制作为 OFRecord 数据集

OFRecord 数据格式加载与准备 OFRecord 数据集 中,我们分别学习了 OFRecord 数据格式,以及如何将其它数据集转为 OFRecord 数据集并使用。

本文,我们将介绍如何将图片文件制作为 OFRecord 数据集,并提供了相关的制作脚本,方便用户直接使用或者在此基础上修改。内容包括:

  • 制作基于 MNIST 手写数字数据集的 OFRecord 数据集
  • OFRecord 的编解码方式
  • 在自制的 OFRecord 数据集上进行训练

用图片文件制作 OFRecord 文件

我们使用 MNIST 数据集中的图片文件来制作一个 OFRecord 格式文件。

作为示例,我们仅使用了50张图片,相关脚本和数据集的下载地址为 img2ofrecord

  • 下载相关压缩包并解压
  1. $ wget https://oneflow-static.oss-cn-beijing.aliyuncs.com/oneflow-tutorial-attachments/img2ofrecord.zip
  2. $ unzip img2ofrecord.zip
  • 进入到对应目录,并运行 OFRecord 制作脚本 img2ofrecord.py
  1. $ cd ./img_to_ofrecord
  2. $ python img2ofrecord.py --part_num=5 --save_dir=./dataset/ --img_format=.png --image_root=./images/train_set/
  • 脚本运行过程中,将输出以下内容
  1. The image root is: ./images/train_set/
  2. The amount of OFRecord data part is: 5
  3. The directory of Labels is: ./images/train_label/label.txt
  4. The image format is: .png
  5. The OFRecord save directory is: ./dataset/
  6. Start Processing......
  7. ./images/train_set/00000030_3.png feature saved
  8. ./images/train_set/00000034_0.png feature saved
  9. ./images/train_set/00000026_4.png feature saved
  10. ./images/train_set/00000043_9.png feature saved
  11. ......
  12. Process image successfully !!!

至此我们 OFRecord 文件制作完毕,并保存在 ./dataset 目录下

代码解读

整个代码目录构造如下

  1. img_to_ofrecord
  2. ├── images
  3. ├── train_set
  4. ├── 00000000_5.png
  5. ├── 00000001_0.png
  6. ├── 00000002_4.png
  7. ......
  8. ├── train_label
  9. ├── label.txt
  10. ├── img2ofrecord.py
  11. ├── lenet_train.py
  • images 目录存放原始示例训练数据集以及标签文件

我们的标签文件是以 json 格式存储的,格式如下:

  1. {"00000030_3.png": 3}
  2. {"00000034_0.png": 0}
  3. {"00000026_4.png": 4}
  4. {"00000043_9.png": 9}
  5. {"00000047_5.png": 5}
  6. {"00000003_1.png": 1}
  7. ......
  • img2ofrecord.py 脚本将 MNIST 图片转换成 OFRecord 数据集
  • lenet_train.py 脚本则读取我们制作好的 OFRecord 数据集,并使用 LeNet 模型进行训练。

img2ofrecord.py 的命令行选项如下:

  • image_root 指定图片的根目录路径
  • part_num 指定生成 OFRecord 文件个数,如果该数目大于总图片数目,会报错
  • label_dir 指定标签的目录路径
  • img_format 指定图片的格式
  • save_dir 指定 OFRecord 文件保存的目录

脚本的编码流程

与 OFRecord 文件编码的相关逻辑也在 img2ofrecord.py 内,其编码流程如下:

首先,对读取进来的图片数据进行编码

  1. def encode_img_file(filename, ext=".jpg"):
  2. img = cv2.imread(filename)
  3. encoded_data = cv2.imencode(ext, img)[1]
  4. return encoded_data.tostring()

这里的 ext 是图片编码格式,目前,OneFlow 图片编解码支持的格式与 OpenCV 的一致,可参见 cv::ImwriteFlags,包括:

然后,转化成 Feature 的形式,进行序列化,并将数据长度写入到文件中

  1. def ndarray2ofrecords(dsfile, dataname, encoded_data, labelname, encoded_label):
  2. topack = {dataname: bytes_feature(encoded_data),
  3. labelname: int32_feature(encoded_label)}
  4. ofrecord_features = ofrecord.OFRecord(feature=topack)
  5. serilizedBytes = ofrecord_features.SerializeToString()
  6. length = ofrecord_features.ByteSize()
  7. dsfile.write(struct.pack("q", length))
  8. dsfile.write(serilizedBytes)

使用自制的 OFRecord 数据集进行训练

我们运行目录下的 lenet_train.py,它将读取我们刚制作好的 OFRecord 数据集,在 Lenet 模型上进行训练

该训练脚本输出如下:

  1. [6.778578]
  2. [2.0212684]
  3. [1.3814741]
  4. [0.47514156]
  5. [0.13277876]
  6. [0.16388433]
  7. [0.03788032]
  8. [0.01225162]
  9. ......

至此,我们成功完成了数据集制作、读取与训练整个流程。

为正常使用来必力评论功能请激活JavaScript