Dataset Grouper - 可扩展的组结构学习数据集管道

Dataset Grouper是一个用于创建、编写和迭代具有组级结构的数据集的库。它主要用于创建大规模数据集，用于联邦学习研究。

安装

我们建议通过PyPI进行安装。请检查PyPI页面，以获取最新的版本要求，包括Python版本要求。

pip install --upgrade pip
pip install dataset-grouper

入门

下面是一个简单的入门示例，将MNIST数据集按标签分配给10个客户端。

首先，我们导入一些必要的包，包括Apache Beam，它将用于运行Dataset Grouper的管道。

import apache_beam as beam
import dataset_grouper as dsgp
import tensorflow_datasets as tfds

接下来，我们下载并准备MNIST数据集。

dataset_builder = tfds.builder('mnist')
dataset_builder.download_and_prepare(...)

现在我们编写一个函数，为每个MNIST示例分配一个客户端标识符（通常是一个bytes对象）。在这种情况下，我们将根据标签划分示例，但您也可以使用更多有趣的划分函数。

def label_partition(x):
  label = x['label'].numpy()
  return str(label).encode('utf-8')

最后，我们构建一个管道，根据此函数划分MNIST，并使用Beam的Direct Runner运行它。

mnist_pipeline = dsgp.tfds_to_tfrecords(
    dataset_builder=dataset_builder,
    split='test',
    get_key_fn=label_partition,
    file_path_prefix=...
)
with beam.Pipeline() as root:
  mnist_pipeline(root)

这将保存一个根据标签划分的MNIST版本，并将其保存为TFRecord格式。我们还可以将其加载以迭代客户端数据集。

partitioned_dataset = dsgp.PartitionedDataset(
  file_pattern=...,
  tfds_features='mnist')

for group_dataset in partitioned_dataset.build_group_stream():
  pass

通常，PartitionedDataset.build_group_stream()是一个tf.data.Dataset，它会生成数据集，每个数据集包含一个组中所有示例。如果您想使用这些数据集与NumPy一起使用，您可以直接这样做：

group_dataset_numpy = group_dataset.as_numpy_iterator()

还有什么吗？

上面的例子主要是为了教育目的。MNIST是一个相对较小的数据集，通常可以完全放入内存中。对于更有趣的例子，请查看示例文件夹。

数据集分组器更适合大型数据集，特别是那些不适合内存的数据集。对于这些数据集，我们建议使用更复杂的Beam runners，以便以分布式方式划分数据。

免责声明

这不是一个官方支持的产品。

这是一个下载和准备公共数据集的实用程序库。我们不托管或分发这些数据集，也不保证其质量或公平性，也不声称您有权使用数据集。您有责任确定您是否有权根据数据集的许可证使用数据集。

如果您想了解更多关于负责任的AI实践，请参阅Google AI的负责任的AI实践。

数据集分组器是Apache 2.0许可的。请参阅LICENSE文件。