在机器学习的背景下,分片具体指什么(一个更通用的问题在这里讨论),以及如何在Tensorflow中实现分片?
我们所说的分片是什么,为什么在机器学习的数据管道中需要分片?
回答:
在Tensorflow中 – 在Dataset
中,函数shard()
创建了一个只包含该数据集1/num_shards的Dataset。分片是确定性的。由A.shard(n, i)生成的Dataset将包含所有满足A的索引mod n = i的元素。
A = tf.data.Dataset.range(10)B = A.shard(num_shards=3, index=0)list(B.as_numpy_iterator())[0,3,6,9]C = A.shard(num_shards=3, index=1)list(C.as_numpy_iterator())[1,4,7]D = A.shard(num_shards=3, index=2)list(D.as_numpy_iterator())[2,5,8]
重要注意事项:确保在使用任何随机化操作符(如shuffle)之前进行分片。
通常,最好在数据集管道的早期使用分片操作符。例如,在从一组TFRecord文件读取时,在将数据集转换为输入样本之前进行分片。这可以避免每个工作节点都读取每个文件。以下是一个在完整管道中高效分片策略的示例:
在多个工作节点上自动分片一个数据集意味着每个工作节点被分配了整个数据集的一个子集(如果设置了正确的tf.data.experimental.AutoShardPolicy)。
这样可以确保在每个步骤中,每个工作节点处理一个非重叠的数据集元素的全局批次大小。
设置自动分片选项的示例
dataset = tf.data.Dataset.from_tensors(([1.],[1.])).repeat(64).batch(16)options = tf.data.Options()options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATAdataset = dataset.with_options(options)
在使用ParameterServerStrategy
的多工作节点训练中没有自动分片。
自动分片选项包括:
-
AUTO:这是默认选项,意味着会尝试按文件分片。如果未检测到基于文件的数据集,则尝试按文件分片会失败。
-
FILE:如果您想在所有工作节点上分片输入文件,应该使用此选项。如果输入文件的数量远大于工作节点的数量,并且文件中的数据分布均匀,则应使用此选项。例如,让我们将2个文件分配给2个工作节点,每个节点有1个副本。文件1包含[0, 1, 2, 3, 4, 5],文件2包含[6, 7, 8, 9, 10, 11]。假设同步的总副本数为2,全局批次大小为4。
工作节点0:批次1 = 副本1:[0, 1]批次2 = 副本1:[2, 3]批次3 = 副本1:[4]批次4 = 副本1:[5]工作节点1:批次1 = 副本2:[6, 7]批次2 = 副本2:[8, 9]批次3 = 副本2:[10]批次4 = 副本2:[11]
-
DATA:这将在所有工作节点上自动分片元素。每个工作节点将读取整个数据集,并只处理分配给它的分片。所有其他分片将被丢弃。这通常在输入文件的数量少于工作节点的数量,并且您希望在所有工作节点上更好地分片数据时使用。其缺点是每个工作节点将读取整个数据集。例如,让我们将1个文件分配给2个工作节点。文件1包含[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]。假设同步的总副本数为2。
工作节点0:批次1 = 副本1:[0, 1]批次2 = 副本1:[4, 5]批次3 = 副本1:[8, 9]工作节点1:批次1 = 副本2:[2, 3]批次2 = 副本2:[6, 7]批次3 = 副本2:[10, 11]
-
OFF:如果您关闭自动分片,每个工作节点将处理所有数据。例如,让我们将1个文件分配给2个工作节点。文件1包含[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]。假设同步的总副本数为2。那么每个工作节点将看到以下分配:
工作节点0:批次1 = 副本1:[0, 1]批次2 = 副本1:[2, 3]批次3 = 副本1:[4, 5]批次4 = 副本1:[6, 7]批次5 = 副本1:[8, 9]批次6 = 副本1:[10, 11]
工作节点1:批次1 = 副本2:[0, 1]批次2 = 副本2:[2, 3]批次3 = 副本2:[4, 5]批次4 = 副本2:[6, 7]批次5 = 副本2:[8, 9]批次6 = 副本2:[10, 11]