ray.train.torch.prepare_数据加载器#

ray.train.torch.prepare_data_loader(data_loader: torch.utils.data.DataLoader, add_dist_sampler: bool = True, move_to_device: bool = True, auto_transfer: bool = True) → torch.utils.data.DataLoader[源代码]#

为分布式执行准备 DataLoader。

这允许你使用相同的代码，无论工作者的数量或使用的设备类型（CPU、GPU）如何。

备注

如果训练工作者的数量大于1，此方法会向 DataLoader 添加一个 DistributedSampler。如果原始 DataLoader 启用了洗牌功能，那么 shuffle=True 也会传递给 DistributedSampler 构造函数。原始 DataLoader 上的 shuffle=False 也意味着在采样器上禁用了洗牌功能。

使用超过1个工作进程时，在每个epoch开始之前调用 DistributedSampler.set_epoch 方法，然后再创建 DataLoader 迭代器是必要的，以确保在多个epoch中正确地进行洗牌。否则，将始终使用相同的顺序。参见：https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler # noqa: E501

示例：

参数:

  • data_loader (torch.utils.data.DataLoader) – 用于准备的 DataLoader。

  • add_dist_sampler – 是否将 DistributedSampler 添加到提供的 DataLoader 中。

  • move_to_device – 如果设置，自动将数据加载器返回的数据移动到正确的设备。

  • auto_transfer – 如果设置且设备为GPU，则会创建另一个CUDA流来自动将数据从主机（CPU）内存复制到设备（GPU）内存（默认的CUDA流仍然运行训练过程）。如果设备为CPU，则无论设置如何都将禁用此功能。如果 move_to_device 为False，则此配置将被忽略。