python - Tensorflow:如何手动对数据集进行分片

Question

我正在使用 MirroredStrategy 执行多 GPU 训练，但它似乎没有正确地对数据进行分片。如何手动对数据进行分片？

我知道我可以对 tf.data 数据集使用 shard 方法，但为此我需要访问工作人员 ID，但我无法弄清楚如何得到它。如何访问工作人员 ID？

Answer 1

MirroredStrategy 在单个工作线程上运行(对于多个工作线程，有 MultiWorkerMirroredStrategy )。由于仅在一个工作线程上运行，MirroredStrategy 运行单个 Dataset 管道，无需任何数据分片。在每一步中，MirroredStrategy 都会为每个工作线程请求一个数据集元素。