TensorFlow Dataset save未生成多分片，用法存疑

问题核心原因 #

tf.data.Dataset.save默认仅生成1个分片，因为该方法的num_shards参数默认值为1，且默认写入线程数为1。官方文档中提到的“轮询方式分配数据”是当指定多个分片（num_shards > 1）时的分配逻辑，并非自动拆分出多个分片。

生成多分片的解决方法 #

只需在调用save时显式设置num_shards参数，指定想要的分片数量即可。示例代码修改如下：

import tempfile
import tensorflow as tf
import os

path = os.path.join(tempfile.gettempdir(), "saved_data")
dataset = tf.data.Dataset.range(10**8)

# 指定生成4个分片，可根据需求调整数量
dataset.save(path, num_shards=4)

执行后会生成形如data-00000-of-00004、data-00001-of-00004的多个分片文件，数据会按轮询策略自动分配到各个分片。

额外说明 #

• 若需要自定义数据分配到分片的逻辑，可以通过shard_func参数传入自定义函数，函数接收元素的索引和值，返回该元素所属的分片ID（需在0到num_shards-1范围内）。
• 分片数量的选择需结合存储系统性能和后续加载时的并行需求，过多小分片可能增加IO开销，需合理权衡。

内容的提问来源于stack exchange，提问作者melec petit-pierre