Dpex——用户无感知分布式数据预处理组件
一、前言
随着GPU与CPU的算力差距越来越大以及模型训练时的预处理Pipeline变得越来越复杂，CPU部分的数据预处理已经逐渐成为了模型训练的瓶颈所在，这导致单机的GPU配置的提升并不能带来期望的线性加速。预处理性能瓶颈的本质在于每个GPU能够使用的CPU算力受限， 为了解决这个问题NVIDIA提出了scale up的方案——GPU数据预处理库DALI，Tensorflow给出了scale out的方案——分布式数据预处理组件DataService,而在这里我们给出Pytorch生态中的scale out方案——分布式数据预处理组件Dpex。
二、架构介绍（介绍Pytorch DataLoader本身的架构以及DistDataLoader的架构）
Dpex的采用了和Pytorch的DataLoader同样的架构设计并借助Ray将数据预处理任务调度至其他机器节点进行计算。 
三、使用示例
不仅在设计上，Dpex的实现上也完全兼容Pytorch的DataLoader。当并行数据预处理时，若设置distribute_mode为True则DpexDataLoader使用 _RayDataLoaderIter实现分布式数据预处理，当设置为False时DpexDataLoader退回到使用Pytorch本身的_MultiProcessingDataLoaderIter 实现并行数据预处理与加载。在Pytorch训练中使用Dpex非常的简单，只需要将使用到Pytorch的DataLoader的地方替换为Dpex中的DpexDataLoader即可，当你的训练机器本身为Ray集群中的一个节点时，设置 distribute_mode=True可以启用分布式数据预处理。在下面我们给出单卡训练，使用DataParallel进行多卡训练以及使用DDP进行多卡训练时使用Dpex的示例，具体可参考测试文件。
class DpexDataLoader(torch.utils.data.DataLoader): def init(self, dataset: Dataset[T_co], distribute_mode: Optional[bool] = False, head_address="auto", batch_size: Optional[int] = 1, shuffle: bool = False, sampler: Optional[Sampler[int]] = None, batch_sampler: Optional[Sampler[Sequence[int]]] = None, num_workers: int = 0, collate_fn: Optional[_collate_fn_t] = None, pin_memory: bool = False, drop_last: bool = False, timeout: float = 0, worker_init_fn: Optional[_worker_init_fn_t] = None, multiprocessing_context=None, generator=None, *, prefetch_factor: int = 2):
3.1 单卡训练
如下我们给出单卡训练时使用DpexDataLoader的示例代码，具体代码细节参见测试代码文件.
from torchvision import datasets
from torchvision.transforms import ToTensor
from Dpex import dataloader
training_data = datasets.FashionMNIST(
root="data",
train=True,
download=True,
transform=ToTensor()
)
# use DpexDataLoader
train_loader = dataloader.DpexDataLoader(training_data, distribute_mode=True, num_workers=10, batch_size=100, shuffle=True)
for epoch in range(3):
for index, (image, label) in enumerate(train_loader):
# your train process
pass
3.2 基于DataParallel的多卡训练
如下我们给出使用DataParallel并行训练时使用DpexDataLoader的示例代码，具体代码细节参见测试代码文件.
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import Dataset
from Dpex import dataloader
class MyOwnDataset(Dataset):
pass
# use DpexDataLoader
data_loader = dataloader.DpexDataLoader(dataset=RandomDataset(input_size, data_size),
distribute_mode=True, batch_size=batch_size, shuffle=True, num_workers=10)
class Model(nn.Module):
pass
model = Model()
if torch.cuda.is_available():
model.cuda()
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
for data in data_loader:
# train your own model
pass
3.3 基于DDP的多卡训练
如下我们给出使用DDP并行训练时使用DpexDataLoader的示例代码，具体代码细节参见测试代码文件.
import torch
import torch.nn as nn
from torch.utils.data import Dataset
from Dpex.dataloader import DpexDataLoader
from torch.utils.data.distributed import DistributedSampler
# start command: CUDA_VISIBLE_DEVICES=1,6,7 python -m torch.distributed.launch --nproc_per_node=2 pytorch_ddp.py
# 1) 初始化
torch.distributed.init_process_group(backend="nccl")
input_size = 5
output_size = 2
batch_size = 1
data_size = 90000
# 2） 配置每个进程的gpu
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
class RandomDataset(Dataset):
def __init__(self, size, length):
self.len = length
self.data = torch.randn(length, size)
def __getitem__(self, index):
return self.data[index]
def __len__(self):
return self.len
dataset = RandomDataset(input_size, data_size)
# 3）使用DistributedSampler
rand_loader = DpexDataLoader(dataset=dataset, distribute_mode=True, batch_size=batch_size, sampler=DistributedSampler(dataset), num_workers=10)
class Model(nn.Module):
def __init__(self, input_size, output_size):
super(Model, self).__init__()
self.fc = nn.Linear(input_size, output_size)
def forward(self, input):
output = self.fc(input)
print(" In Model: input size", input.size(),
"output size", output.size())
return output
model = Model(input_size, output_size)
# 4) 封装之前要把模型移到对应的gpu
model.to(device)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# 5) 封装
model = torch.nn.parallel.DistributedDataParallel(model,
device_ids=[local_rank],
output_device=local_rank)
for data in rand_loader:
if torch.cuda.is_available():
input_var = data
else:
input_var = data
output = model(input_var)
print("Outside: input size", input_var.size(), "output_size", output.size())
四、Benchmark
在接下来的Benchamark中我们核心展示两个部分的内容:
- DpexDataLoader对于模型训练精度的影响
- DpexDataLoader对于模型训练速度的影响
Dpex只是将单机数据预处理水平扩展到了多机以借助更多的算力来加速数据预处理而不改变数据本身的加载和与处理方式，所以本身对模型的精度不会有负面影响。对于数据预处理较重的情况
4.1 模型精度Benchmark
我们在FashionMNIST数据集上进行模型训练精度的Benchmark实验，具体代码细节见测试文件
| Accuracy(%) | Loss | GPU Settings | DpexDataLoader(distributed_mode=?) | Epoch | Learning rate | Batch size |
|---|---|---|---|---|---|---|
| 90.65 | 0.137 | Single GPU | True | 40 | 0.001 | 100 |
| 91.09 | 0.112 | Single GPU | False | 40 | 0.001 | 100 |
| 90.67 | 0.016 | DataParallel | True | 40 | 0.001 | 100 |
| 90.32 | 0.008 | DataParallel | False | 40 | 0.001 | 100 |
| 88.98 | 0.034 | DDP | True | 40 | 0.001 | 100 |
| 89.84 | 0.030 | DDP | False | 40 | 0.001 | 100 |
4.2 训练速度Benchmark
五、环境依赖:
Dpex借助Ray完成任务的跨机调度，所以若希望使用分布式数据预处理首先需要将你的训练机器构建成Ray的集群。Ray的集群构建细节具体参考Ray的相关文档
157 Dec 26, 2022
575 Jan 8, 2023
26 Dec 13, 2022
2 Mar 30, 2022
3 Jan 8, 2022
5 Jan 15, 2022
14.5k Jan 8, 2023
163 Dec 26, 2022
208 Dec 20, 2022
5 Mar 25, 2022
1 Nov 19, 2021
166 Dec 27, 2022
7 Dec 27, 2022
2 Dec 26, 2021
9 Sep 18, 2022
6.9k Jan 3, 2023
5 Sep 29, 2021
32 Dec 25, 2022