如何使用张量流数据集碎片使每个工人具有相同的铲斗长度以进行分布式培训

Question

下面是我的代码

train_dataset = tf.data.Dataset.list_files(file_pattern=file_pattern, shuffle=True).shard(
    shard_nums, shard_index).repeat(500).shuffle(1000).interleave(
    tf.data.TFRecordDataset, cycle_length=4, block_length=64).map(
    parser_example, num_parallel_calls=64
).apply(
    tf.contrib.data.bucket_by_sequence_length(
        element_length_func=element_length_fn,
        bucket_boundaries=[i for i in range(0, 1000, 50)],
        bucket_batch_sizes=[64 for _ in range(0, 1000, 50)] + [64],
        padded_shapes=([None, hparams.num_freq], [None, hparams.num_mels], [None], [], []),
        padding_values=(0., 0., 0, 0, 0),
        pad_to_bucket_boundary=False
    )).prefetch(1)

我想使用张量流进行分布式训练，但是我想给每个工作人员存储桶长度几乎相同，但是，数据集api shard只能为工作人员制作元素，我该怎么做？

一台服务器和四个woker，等级0、1、2、3 Rank 3, input_shape: (64, 345, 513) Rank 1, input_shape: (64, 145, 513) Rank 0, input_shape: (64, 145, 513) Rank 2, input_shape: (64, 145, 513)

因此您可以看到3级的长度为345，其他的长度为145，花费的时间取决于345，这会浪费很多时间。希望有人有解决方案。