从 C+ 生成 TFRecord 格式数据

Question

我正在尝试使用 TFRecord format 从 C++ 记录数据，然后在 python 中使用它来提供 TensorFlow 模型。

TLDR；简单地将 proto 消息序列化为流不满足 Python TFRecordDataset class 的 .tfrecord 格式要求。在 C++ 中（在 TensorFlow 或 Google Protobuf 库中）是否有等效的 Python TfRecordWriter 来生成适当的 .tfrecord 数据？

详情：

简化的 C++ 代码如下所示：

tensorflow::Example sample;
sample.mutable_features()->mutable_feature()->operator[]("a").mutable_float_list()->add_value(1.0);

std::ofstream out;
out.open("cpp_example.tfrecord", std::ios::out | std::ios::binary);
sample.SerializeToOstream(&out);

在 Python 中，我尝试使用 TFRecordDataset 创建 TensorFlow 数据，但显然它需要 .tfrecord 文件中的额外 header/footer 信息（而不是简单的列表序列化原型消息）：

import tensorflow as tf
tfrecord_dataset = tf.data.TFRecordDataset(filenames="cpp_example.tfrecord")
next(tfrecord_dataset.as_numpy_iterator())

输出：

tensorflow.python.framework.errors_impl.DataLossError: corrupted record at 0 [Op:IteratorGetNext]

请注意，记录的二进制文件没有任何问题，因为以下代码打印出有效输出：

import tensorflow as tf
p = open("cpp_example.tfrecord", "rb")
example = tf.train.Example.FromString(p.read())

输出：

features {
  feature {
    key: "a"
    value {
      float_list {
        value: 1.0
      }
    }
  }
}

通过分析我的 C++ 示例生成的二进制输出，以及使用 Python TfRecordWriter 生成的输出，我观察到内容中有额外的页眉和页脚字节。不幸的是，这些额外的字节代表的是一个实现细节（可能是压缩类型和一些额外的信息），我无法比 python 库中的某些 class 更深入地跟踪它，这些库只是从 _pywrap_tfe.so.

有 this advice 说 .tfrecord 只是一个普通的 google protobuf 数据。可能是我不知道在哪里可以找到 protobuf 数据编写器（期望将 proto 消息序列化到输出流中）？

Answer 1

事实证明，TensorFlow C++ 库的 tensorflow::io::RecordWriter class 可以完成这项工作。

#include <tensorflow/core/lib/io/record_writer.h>

#include <tensorflow/core/platform/default/posix_file_system.h>
#include <tensorflow/core/example/example.pb.h>

// ...

// Create WritableFile and instantiate RecordWriter.
tensorflow::PosixFileSystem posixFileSystem;
std::unique_ptr<tensorflow::WritableFile> writableFile;

posixFileSystem.NewWritableFile("cpp_example.tfrecord", &writableFile);

tensorflow::io::RecordWriter recordWriter(mWritableFile.get(), tensorflow::io::RecordWriterOptions::CreateRecordWriterOptions(""));

// ...
tensorflow::Example sample;

// ...

// Serialize proto message into a buffer and record in tfrecord format.
std::string buffer;
sample.SerializeToString(&buffer);
recordWriter.WriteRecord(buffer);

如果从 TFRecord documentation 中的某处引用此 class 将会很有帮助。

从 C+ 生成 TFRecord 格式数据

Generating TFRecord format data from C+

c++

python

protocol-buffers

tensorflow

tfrecord