如何将 Tensorflow 数据集保存到文件中？

Question

在 SO 上至少还有两个这样的问题，但没有一个得到回答。

我有以下形式的数据集：

<TensorSliceDataset shapes: ((512,), (512,), (512,), ()), types: (tf.int32, tf.int32, tf.int32, tf.int32)>

和另一种形式：

<BatchDataset shapes: ((None, 512), (None, 512), (None, 512), (None,)), types: (tf.int32, tf.int32, tf.int32, tf.int32)>

我看了又看，但找不到将这些数据集保存到以后可以加载的文件的代码。我得到的最接近的是TensorFlow 文档中的这个页面，它建议使用序列化张量tf.io.serialize_tensor，然后使用tf.data.experimental.TFRecordWriter.

但是，当我使用代码尝试此操作时：

dataset.map(tf.io.serialize_tensor)
writer = tf.data.experimental.TFRecordWriter('mydata.tfrecord')
writer.write(dataset)

我在第一行收到错误：

类型错误：serialize_tensor() 需要 1 到 2 个位置参数，但给出了 4 个

我如何修改上述（或做其他事情）以实现我的目标？

Answer 1

GitHUb 上发生了一个事件，看来 TF 2.3 中有一个新功能可以写入磁盘：

https://www.tensorflow.org/api_docs/python/tf/data/experimental/save https://www.tensorflow.org/api_docs/python/tf/data/experimental/load

我还没有测试过这个功能，但它似乎正在做你想要的。

Answer 2

TFRecordWriter似乎是最方便的选择，但不幸的是它只能编写每个元素一个张量的数据集。您可以使用以下几种解决方法。首先，由于您所有的张量都具有相同的类型和相似的形状，您可以将它们全部连接成一个，然后在加载时将它们拆分回来：

import tensorflow as tf

# Write
a = tf.zeros((100, 512), tf.int32)
ds = tf.data.Dataset.from_tensor_slices((a, a, a, a[:, 0]))
print(ds)
# <TensorSliceDataset shapes: ((512,), (512,), (512,), ()), types: (tf.int32, tf.int32, tf.int32, tf.int32)>
def write_map_fn(x1, x2, x3, x4):
    return tf.io.serialize_tensor(tf.concat([x1, x2, x3, tf.expand_dims(x4, -1)], -1))
ds = ds.map(write_map_fn)
writer = tf.data.experimental.TFRecordWriter('mydata.tfrecord')
writer.write(ds)

# Read
def read_map_fn(x):
    xp = tf.io.parse_tensor(x, tf.int32)
    # Optionally set shape
    xp.set_shape([1537])  # Do `xp.set_shape([None, 1537])` if using batches
    # Use `x[:, :512], ...` if using batches
    return xp[:512], xp[512:1024], xp[1024:1536], xp[-1]
ds = tf.data.TFRecordDataset('mydata.tfrecord').map(read_map_fn)
print(ds)
# <MapDataset shapes: ((512,), (512,), (512,), ()), types: (tf.int32, tf.int32, tf.int32, tf.int32)>

但是，更一般地，您可以简单地为每个张量创建一个单独的文件，然后将它们全部读取：

import tensorflow as tf

# Write
a = tf.zeros((100, 512), tf.int32)
ds = tf.data.Dataset.from_tensor_slices((a, a, a, a[:, 0]))
for i, _ in enumerate(ds.element_spec):
    ds_i = ds.map(lambda *args: args[i]).map(tf.io.serialize_tensor)
    writer = tf.data.experimental.TFRecordWriter(f'mydata.{i}.tfrecord')
    writer.write(ds_i)

# Read
NUM_PARTS = 4
parts = []
def read_map_fn(x):
    return tf.io.parse_tensor(x, tf.int32)
for i in range(NUM_PARTS):
    parts.append(tf.data.TFRecordDataset(f'mydata.{i}.tfrecord').map(read_map_fn))
ds = tf.data.Dataset.zip(tuple(parts))
print(ds)
# <ZipDataset shapes: (<unknown>, <unknown>, <unknown>, <unknown>), types: (tf.int32, tf.int32, tf.int32, tf.int32)>

可以将整个数据集放在一个文件中，每个元素有多个单独的张量，即作为包含tf.train.Examples的 TFRecords 文件，但我不知道是否有办法在 TensorFlow 中创建它们，也就是说，无需将数据从数据集中提取到 Python 中，然后将其写入记录文件。

Answer 3

补充Yoan的答案：

tf.experimental.save() 和 load() API 运行良好。您还需要手动将 ds.element_spec 保存到磁盘，以便稍后/在不同的上下文中加载（）。

酸洗对我来说效果很好：

1- 保存：

tf.data.experimental.save(
    ds, tf_data_path, compression='GZIP'
)
with open(tf_data_path + '/element_spec', 'wb') as out_:  # also save the element_spec to disk for future loading
    pickle.dump(ds.element_spec, out_)

2-为了加载，您需要包含 tf 分片的文件夹路径和我们手动腌制的 element_spec

with open(tf_data_path + '/element_spec', 'rb') as in_:
    es = pickle.load(in_)

loaded = tf.data.experimental.load(
    tf_data_path, es, compression='GZIP'
)

Answer 4

Tensorflow 2.10 将save方法从tf.data.experimental移至tf.data.Dataset。与load方法一起使用，这是保存和加载模型的最简单方法。

Tensorflow 2.6 引入了快照方法（以前是“实验性”功能）。 Tensorflow RFC-193详细介绍了该功能的动机和细节。

来自文档：

快照 API 允许用户透明地将预处理管道的输出保存到磁盘，并在不同的训练运行中具体化预处理数据。

该 API 可以整合重复的预处理步骤，并允许重复使用已处理的数据，权衡磁盘存储和网络带宽，以释放更有价值的 CPU 资源和加速器计算时间。