OFRecord 图片文件制数据集

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

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

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

用图片文件制作 OFRecord 文件

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

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

• 下载相关压缩包并解压

$ wget https://oneflow-static.oss-cn-beijing.aliyuncs.com/oneflow-tutorial-attachments/img2ofrecord.zip

$ unzip img2ofrecord.zip

• 进入到对应目录，并运行 OFRecord 制作脚本 img2ofrecord.py

$ cd ./img_to_ofrecord

$ python img2ofrecord.py --part_num=5 --save_dir=./dataset/ --img_format=.png --image_root=./images/train_set/

• 脚本运行过程中，将输出以下内容

The image root is:  ./images/train_set/

The amount of OFRecord data part is:  5

The directory of Labels is:  ./images/train_label/label.txt

The image format is:  .png

The OFRecord save directory is:  ./dataset/

Start Processing……

./images/train_set/00000030_3.png feature saved

./images/train_set/00000034_0.png feature saved

./images/train_set/00000026_4.png feature saved

./images/train_set/00000043_9.png feature saved

……

Process image successfully !!!

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

代码解读

整个代码目录构造如下

img_to_ofrecord

├── images

├── train_set

├── 00000000_5.png

├── 00000001_0.png

├── 00000002_4.png

……

├── train_label

├── label.txt

├── img2ofrecord.py

├── lenet_train.py

• images 目录存放原始示例训练数据集以及标签文件

的标签文件是以 json 格式存储的，格式如下：

{"00000030_3.png": 3}

{"00000034_0.png": 0}

{"00000026_4.png": 4}

{"00000043_9.png": 9}

{"00000047_5.png": 5}

{"00000003_1.png": 1}

……

• 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 内，其编码流程如下：

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

def encode_img_file(filename, ext=".jpg"):

img = cv2.imread(filename)

encoded_data = cv2.imencode(ext, img)[1]

return encoded_data.tostring()

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

• JPEG，一种最常见的有损编码格式，可参考 JPEG
• PNG，一种常见的无损位图编码格式，可参考 Portable Network Graphics
• TIFF，一种可扩展的压缩编码格式，可参考 Tagged Image File Format

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

def ndarray2ofrecords(dsfile, dataname, encoded_data, labelname, encoded_label):

topack = {dataname: bytes_feature(encoded_data),

labelname: int32_feature(encoded_label)}

ofrecord_features = ofrecord.OFRecord(feature=topack)

serilizedBytes = ofrecord_features.SerializeToString()

length = ofrecord_features.ByteSize()

dsfile.write(struct.pack("q", length))

dsfile.write(serilizedBytes)

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

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

该训练脚本输出如下：

[6.778578]

[2.0212684]

[1.3814741]

[0.47514156]

[0.13277876]

[0.16388433]

[0.03788032]

[0.01225162]

……

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