对齐PyTorch，一文详解OneFlow的DataLoader实现
撰文 | 赵露阳
在最新的OneFlow v0.5.0版本中，我们增加了许多新特性，比如：
新增动态图特性：OneFlow 默认以动态图模式（eager）运行，与静态图模式（graph）相比，更容易搭建网络、调试和验证算法。
面向对象式的动态图接口 nn.Module，熟悉 PyTorch 的用户可以轻松上手。
“一行代码转换 OneFlow 与 PyTorch 网络”：与 PyTorch 对齐的算子数目增加至200+。在 ResNet50、AlexNet 等 十几个常用网络 上已通过 import oneflow as torch 和 import torch as flow 验证。注意：此特性是为方便用户由 PyTorch 迁移至 OneFlow 而设计，并不是承诺完全兼容 PyTorch。
面向对象式的静态图接口：新增面向对象的静态图接口 nn.Graph。保留了 OneFlow 静态图性能优势的同时，让静态图的编程门槛与动态图接近，期待更多的算法工程师把 OneFlow 的高性能优势玩起来。这是一个用 nn.Graph 搭建 ResNet50 示例
易用高效的分布式训练：分布式训练是大势所趋，OneFlow 本版本新增的 Consistent Tensor，让用户可以像操作单机单卡一样，操作整个集群，并立即看到效果。新增的 launch 模块、DDP 模块 配合 OneFlow 的一致性视角 让用户轻松启动分布式训练，无论是 数据并行、模型并行、还是流水并行，OneFlow 均原生支持，易用高效。
其中，最重要的新特性之一，就是OneFlow的动态图做到了几乎和PyTorch一致，从Tensor、nn.Module、到autograd、functional api等，其中也包括和torch几乎对齐的DataLoader/Dataset设计，笔者有幸开发了OneFlow中的这一模块。
https://github.com/Oneflow-Inc/oneflow/pull/5406
https://github.com/Oneflow-Inc/oneflow/pull/5500
https://github.com/Oneflow-Inc/oneflow/pull/5644
https://github.com/Oneflow-Inc/oneflow/pull/6280本文将对OneFlow/PyTorch中的DataLoader原理、工作流程进行梳理：
dataloader简介
dataloader原理
dataloader工作流程
multiprocessing dataloader工作原理
1
简介
简单来说，DataLoader是深度学习中必不可少的，用于处理Dataset产生每个iter过程中批量数据和label的一种数据加载器。正如PyTorch文档中的描述：DataLoader，结合了Sampler、Dataset，提供了对某个dataset可迭代的数据集合。DataLoader支持单进程、多进程的加载数据集合。
2
dataloader原理
核心组建
Dataloader
Dataset
Sampler
Fetcher
DataLoader工作原理的简单总结：
1.Dataloader是负责数据加载的核心；DataLoaderIter是具体执行单位。dataloader进入到每一次iter中都会通过DataloaderIter来处理具体的数据加载过程；
2.Dataset是数据集的基类，任何自定义数据集都需要继承它并通过重写getitem方法来定义取数据的方式；
3.Sampler是负责index相关的采样器、每个iter迭代都会通过Sampler生成要采样的数据集的index；
4.Fetcher更像是数据的收集器。根据Sampler产生的batch个index去数据集中fetch对应的数据、并通过相应的collate_fn方法将获取的数据收集打包成最终可用的形式，返回给DataLoader。
使用示例
1.MNIST
下面用PyTorch官方examples的一个简单例子，用MNIST数据集训练分类网络来说明DataLoader的用法：
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
dataset1 = datasets.MNIST('../data', train=True, download=True,
transform=transform)
dataset2 = datasets.MNIST('../data', train=False,
transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1,**train_kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)可以看到，dataset1、dataset2分别是表示数据集的训练集、测试集。在PyTorch中是通过torchvision.datasets.MNIST定义的。MNIST继承自VisionDataset，而VisionDataset则继承自torch.utils.data.Dataset。在MNIST中，实现了数据集最重要的getitem方法，用于根据index取对应数据：
def __getitem__(self, index: int) -> Tuple[Any, Any]:
"""
Args:
index (int): Index
Returns:
tuple: (image, target) where target is index of the target class.
"""
img, target = self.data[index], int(self.targets[index])
# doing this so that it is consistent with all other datasets
# to return a PIL Image
img = Image.fromarray(img.numpy(), mode='L')
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target在OneFlow中，oneflow.utils.data对应torch.utils.data；flowvision对应torchvision，使用方式几乎完全一致。例如：对应MNIST数据集，即可直接通过flowvision.datasets.MNIST使用。
dataset1、dataset2定义完成后，传入分别用于训练、验证的dataloader（train_loader、test_loader）。之后，在train/test的循环中，即可迭代dataloader获取每个iter的数据和label：
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
....2.ImageNet
这里还是用PyTorch官方examples里ImageNet数据集的训练为例：
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
num_workers=args.workers, pin_memory=True, sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)可以看见，大体流程和上面的MNIST差不多：
1.先是构造Dataset，这里为通过datasets.ImageFolder构造。ImageFolder是用于读取/处理以文件夹形式存放的图片数据集：
class ImageFolder(DatasetFolder):
r"""A generic data loader where the images are arranged in this way by default:
.. code-block:: shell
root/dog/xxx.png
root/dog/xxy.png
root/dog/[...]/xxz.png
root/cat/123.png
root/cat/nsdf3.png
root/cat/[...]/asd932_.png
This class inherits from :class:`~vision.datasets.DatasetFolder` so
the same methods can be overridden to customize the dataset.
Args:
root (string): Root directory path.
transform (callable, optional): A function/transform that takes in an PIL image
and returns a transformed version. E.g, ``transforms.RandomCrop``
target_transform (callable, optional): A function/transform that takes in the
target and transforms it.
loader (callable, optional): A function to load an image given its path.
is_valid_file (callable, optional): A function that takes path of an Image file
and check if the file is a valid file (used to check of corrupt files)
Attributes:
classes (list): List of the class names sorted alphabetically.
class_to_idx (dict): Dict with items (class_name, class_index).
imgs (list): List of (image path, class_index) tuples
"""
def __init__(
self,
root: str,
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
loader: Callable[[str], Any] = default_loader,
is_valid_file: Optional[Callable[[str], bool]] = None,
):
super(ImageFolder, self).__init__(
root,
loader,
IMG_EXTENSIONS if is_valid_file is None else None,
transform=transform,
target_transform=target_transform,
is_valid_file=is_valid_file,
)
self.imgs = self.samples可以看到其继承自DatasetFolder、初始化时主要参数有：
root：图片文件夹路径
transform：对经过loader读取到的PIL图片，经过哪些transform处理，如上述的Resize、CenterCrop等
loader：一个用于根据path加载图片的图像加载器，通常默认的loader是PIL
DatasetFolder中实现了Dataset中最重要的getitem方法：
def __getitem__(self, index: int) -> Tuple[Any, Any]:
"""
Args:
index (int): Index
Returns:
tuple: (sample, target) where target is class_index of the target class.
"""
path, target = self.samples[index]
sample = self.loader(path)
if self.transform is not None:
sample = self.transform(sample)
if self.target_transform is not None:
target = self.target_transform(target)
return sample, target通过getitem定义了如何根据index取到相应数据的方式。
2.其次如果是多机分布式训练，则Sampler需要使用专门为分布式训练设计的DistributedSampler类（否则不用特殊设置，用默认的即可）；这里还有个细节，训练集和验证集上，对dataset做了不同的transform，训练集用了RandomResizedCrop、RandomHorizontalFlip；验证集则是Resize、CenterCrop，经过transform后，最终通过ToTensor方法转化成Tensor。
3.构造用于训练、验证的Dataloader(train_loader、val_loader)，后面的使用方式就很简单了，在train/eval的loop中直接使用即可：
for i, (images, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
if args.gpu is not None:
images = images.cuda(args.gpu, non_blocking=True)
if torch.cuda.is_available():
target = target.cuda(args.gpu, non_blocking=True)
.....3
dataloader工作流程
下面结合代码看一下主要流程：
Dataset
任何自定义数据集，必须继承Dataset类并实现_getitem__方法，用于定义根据传入的index获取数据的方式。同时，自定义数据集也可选重写len方法，用于判断数据集的size。
class Dataset(Generic[T_co]):
r"""An abstract class representing a :class:`Dataset`.
All datasets that represent a map from keys to data samples should subclass
it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching a
data sample for a given key. Subclasses could also optionally overwrite
:meth:`__len__`, which is expected to return the size of the dataset by many
:class:`~flow.utils.data.Sampler` implementations and the default options
of :class:`~flow.utils.data.DataLoader`.
.. note::
:class:`~flow.utils.data.DataLoader` by default constructs a index
sampler that yields integral indices. To make it work with a map-style
dataset with non-integral indices/keys, a custom sampler must be provided.
"""
def __getitem__(self, index) -> T_co:
raise NotImplementedError
def __add__(self, other: "Dataset[T_co]") -> "ConcatDataset[T_co]":
return ConcatDataset([self, other])DataLoader
DataLoader是整个数据处理过程的核心。
class DataLoader(Generic[T_co]):
def __init__(
self,
dataset: Dataset[T_co],
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,
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,
persistent_workers: bool = False
):
...
...
# We quote '_BaseDataLoaderIter' since it isn't defined yet and the definition can't be moved up
# since '_BaseDataLoaderIter' references 'DataLoader'.
def __iter__(self) -> "_BaseDataLoaderIter":
# When using a single worker the returned iterator should be
# created everytime to avoid reseting its state
# However, in the case of a multiple workers iterator
# the iterator is only created once in the lifetime of the
# DataLoader object so that workers can be reused
if self.persistent_workers and self.num_workers > 0:
if self._iterator is None:
self._iterator = self._get_iterator()
else:
self._iterator._reset(self)
return self._iterator
else:
return self._get_iterator()
def _get_iterator(self) -> "_BaseDataLoaderIter":
if self.num_workers == 0 or self.num_workers == 1:
return _SingleProcessDataLoaderIter(self)
else:
self.check_worker_number_rationality()
return _MultiProcessingDataLoaderIter(self)DataLoader在每一个iter迭代过程中，最重要的就是通过上面的__iter__方法完成取数据和label。__iter__里通过_get_iterator方法获取相应的DataLoaderIter实例。
在单进程下，即
_SingleProcessDataLoaderIter；多进程下，即
_MultiProcessingDataLoaderIter，他们都继承自_BaseDataLoaderIter
DataLoaderIter
DataLoaderIter负责DataLoader在每个迭代中具体事务的处理。
class _BaseDataLoaderIter(object):
def __init__(self, loader: DataLoader) -> None:
self._dataset = loader.dataset
self._dataset_kind = loader._dataset_kind
self._IterableDataset_len_called = loader._IterableDataset_len_called
self._auto_collation = loader._auto_collation
self._drop_last = loader.drop_last
self._index_sampler = loader._index_sampler
self._num_workers = loader.num_workers
self._prefetch_factor = loader.prefetch_factor
self._pin_memory = False
self._timeout = loader.timeout
self._collate_fn = loader.collate_fn
self._sampler_iter = iter(self._index_sampler)
self._base_seed = flow.tensor([0], dtype=flow.int64).uniform_().numpy().item()
# TODO: flow.empty()
# self._base_seed = flow.empty((), dtype=flow.int64).random_(generator=loader.generator).item()
self._persistent_workers = loader.persistent_workers
self._num_yielded = 0
self._profile_name = "enumerate(DataLoader)#{}.__next__".format(
self.__class__.__name__
)
def __iter__(self) -> "_BaseDataLoaderIter":
return self
def _reset(self, loader, first_iter=False):
self._sampler_iter = iter(self._index_sampler)
self._num_yielded = 0
self._IterableDataset_len_called = loader._IterableDataset_len_called
def _next_index(self):
return next(self._sampler_iter) # may raise StopIteration
def _next_data(self):
raise NotImplementedError
def __next__(self) -> Any:
if self._sampler_iter is None:
self._reset()
data = self._next_data()
self._num_yielded += 1
if (
self._dataset_kind == _DatasetKind.Iterable
and self._IterableDataset_len_called is not None
and self._num_yielded > self._IterableDataset_len_called
):
warn_msg = (
"Length of IterableDataset {} was reported to be {} (when accessing len(dataloader)), but {} "
"samples have been fetched. "
).format(self._dataset, self._IterableDataset_len_called, self._num_yielded)
if self._num_workers > 1:
warn_msg += "Multiprocessing dataloader is not support yet!"
warnings.warn(warn_msg)
return data
def __len__(self) -> int:
return len(self._index_sampler)
def __getstate__(self):
raise NotImplementedError("{} cannot be pickled", self.__class__.__name__)
class _SingleProcessDataLoaderIter(_BaseDataLoaderIter):
def __init__(self, loader):
super(_SingleProcessDataLoaderIter, self).__init__(loader)
assert self._timeout == 0
assert 0 <= self._num_workers <= 1
self._dataset_fetcher = _DatasetKind.create_fetcher(
self._dataset_kind,
self._dataset,
self._auto_collation,
self._collate_fn,
self._drop_last,
)
def _next_data(self):
index = self._next_index() # may raise StopIteration
if self._pin_memory:
raise NotImplementedError("Dataloader pin memory is not support yet!")
return self._dataset_fetcher.fetch(index)在每一个iter迭代时，会调用_BaseDataLoaderIter的__next__方法，进而调用自类实现的_next_data方法获取数据。以_SingleProcessDataLoaderIter为例：
index = self._next_index()通过Sampler获取此次迭代的数据集索引；self._dataset_fetcher.fetch(index)Fetcher根据index索引取相应的数据。
Fetcher
Fetcher作为数据收集器，会根据Sampler产生的batch的index，来从数据集中切分、收集、打包成完整可用的一个batch的数据，并返回给DataLoader使用。
class _MapDatasetFetcher(_BaseDatasetFetcher):
def __init__(self, dataset, auto_collation, collate_fn, drop_last):
super(_MapDatasetFetcher, self).__init__(
dataset, auto_collation, collate_fn, drop_last
)
def fetch(self, possibly_batched_index):
if self.auto_collation:
data = [self.dataset[idx] for idx in possibly_batched_index]
else:
data = self.dataset[possibly_batched_index]
return self.collate_fn(data)Fetcher这里和DataLoaderIter(BaseDataLoaderIter)_类似，_都有一个基类的实现BaseDatasetFetcher。根据不同的数据类型，进入到不同的子类实现中，这里以常用的_MapDatasetFetcher的子类实现为例，看一下Fetcher的主要工作。
可以看见，主要就是：
data = [self.dataset[idx] for idx in possibly_batched_index]return self.collate_fn(data)
1.根据传入的batch个index列表，去dataset中去切分相应的数据，返回的是取出后的batch个数据的列表；
2.根据传入的或自定义的collate_fn方法，收集处理这batch个数据，并打包成训练/验证时可直接使用的Tensor。
4
multiprocessing dataloader工作原理
原理
普通的单进程DataLoader在处理每个iter的数据处理是iter-by-iter且同步的，受制于Python没有实际上的多线程执行，所以单进程的DataLoader通常是比较慢的。多进程DataLoader，即通过Python的multiprocessing开启多个Python的worker进程，譬如开启4个worker进程后，理论上每单位时间可以处理4个iter的数据集，加速数据处理/加载的过程。
单进程DataLoader下，由于数据处理是iter-by-iter的，下一个iter的处理需要等待当前iter完成后才可开始；多进程DataLoader和单进程DataLoader的主要区别就在于可以通过Python的multiprocessing模块，启动多个worker进程加速这个过程。
这里以4进程的DataLoader为例：
DataLoader的主线程将当前iter的任务下发给worker1之后，再下发下一个iter的任务给worker2....直至下发第4个iter的处理任务给worker4。这一步骤主要在dataloader.py的L1024-L1026中实现：
# prime the prefetch loop
for _ in range(self._prefetch_factor * self._num_workers):
self._try_put_index()陆续发送完index后，这4个worker可以并行的工作，陆续完成自己iter的处理任务后，将结果塞入一个Queue队列中，DataLoader的主线程从队列中取数据即可。
具体到每个worker的工作流程，其实和单进程的DataLoader工作流程是类似的，下面主要介绍下多进程和单进程DataLoader的区别，以及多个worker之间是如何协同工作的。
工作流程
_MultiProcessingDataLoaderIter
def _next_data(self):
# DataLoaderIter通过此方法获取每个iter的数据，主要调用_get_data实现
def _get_data(self):
# _get_data方法中，主要通过调用_try_get_data()获取数据
def _try_get_data(self, timeout=_utils.MP_STATUS_CHECK_INTERVAL):
# 从主进程的_data_queue中获取数据
...
try:
data = self._data_queue.get(timeout=timeout)
return (True, data)
except Exception as e:
...
def _process_data(self, data):
# 主要工作即：1.通过_try_put_index()来将下一个iter的index放入一个活跃的worker进程中
# 2.同时标记_rcvd_idx，使其增加1。
self._rcvd_idx += 1
self._try_put_index()
if isinstance(data, ExceptionWrapper):
data.reraise()
return data
def _try_put_index(self):
# 主要工作即遍历所有workers，找到第一个活跃的worker(worker_queue_idx标识)
# 将index和_send_idx信息放入此worker的index_queue中
# 每个worker拥有独立的index_queue，收到index_queue的信息后即开始工作
assert self._tasks_outstanding < self._prefetch_factor * self._num_workers
try:
index = self._next_index()
except StopIteration:
return
for _ in range(self._num_workers): # find the next active worker, if any
worker_queue_idx = next(self._worker_queue_idx_cycle)
if self._workers_status[worker_queue_idx]:
break
else:
# not found (i.e., didn't break)
return
self._index_queues[worker_queue_idx].put((self._send_idx, index))
self._task_info[self._send_idx] = (worker_queue_idx,)
self._tasks_outstanding += 1
self._send_idx += 1_next_data()
⬇️
_get_data() ➡️ _try_get_data()
⬇️
_process_data() ➡️ _try_put_index()
每个worker独立工作，主要代码在oneflow/python/oneflow/utils/data/_utils/worker.py的_worker_loop()方法中：
while watchdog.is_alive():
try:
r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)
except queue.Empty:
continue
if isinstance(r, _ResumeIteration):
# Acknowledge the main process
data_queue.put((r, None))
iteration_end = False
# Recreate the fetcher for worker-reuse policy
fetcher = _DatasetKind.create_fetcher(
dataset_kind, dataset, auto_collation, collate_fn, drop_last
)
continue
elif r is None:
# Received the final signal
assert done_event.is_set() or iteration_end
break
elif done_event.is_set() or iteration_end:
# `done_event` is set. But I haven't received the final signal
# (None) yet. I will keep continuing until get it, and skip the
# processing steps.
continue
idx, index = r
data: Union[_IterableDatasetStopIteration, ExceptionWrapper]
if init_exception is not None:
data = init_exception
init_exception = None
else:
try:
data = fetcher.fetch(index)
except Exception as e:
if (
isinstance(e, StopIteration)
and dataset_kind == _DatasetKind.Iterable
):
data = _IterableDatasetStopIteration(worker_id)
# Set `iteration_end`
# (1) to save future `next(...)` calls, and
# (2) to avoid sending multiple `_IterableDatasetStopIteration`s.
iteration_end = True
else:
# It is important that we don't store exc_info in a variable.
# `ExceptionWrapper` does the correct thing.
# See NOTE [ Python Traceback Reference Cycle Problem ]
data = ExceptionWrapper(
where="in DataLoader worker process {}".format(worker_id)
)
data_queue.put((idx, data))
del data, idx, index, r # save memory
except KeyboardInterrupt:
# Main process will raise KeyboardInterrupt anyways.
pass每个worker在自己的worker loop中，一旦
r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL)获取index_queue中的index数据，就会开始工作：
idx, index = r >> data = fetcher.fetch(index) 这部分内容和之前描述的单进程DataLoader的工作流程没有区别。
当获取到处理完成的数据data后，会将其放入到data loader main线程的data_queue中： data_queue.put((idx, data)) 等待DataLoader主线程从queue中获取结果。
以上即为多进程DataLoader的主要工作流程。
5
结语
本文梳理总结了DataLoader/Dataset，希望能对大家了解OneFlow/PyTorch动态图模式下的DataLoader/Dataset工作原理有所帮助。
对齐PyTorch的DataLoader/Dataset只是第一步，后续仍然面临着效率瓶颈等问题，因为即使使用了multiprocess的DataLoader，在某些情况下，图像解码、Python下调用C++ op执行各种transform时仍可能遭遇性能问题，造成训练过程中GPU打不满/等待CPU数据处理等情况，后续需要考虑更高效的解决方案（如Dali等）。
其他人都在看
OneFlow v0.5.0正式上线
计算机架构史上的一次伟大失败
OneFlow中的错误处理：Maybe
最理想的点到点通信库究竟是怎样的
玩转大模型，是时候展示你真正的技术了
如何实现比PyTorch快6倍的Permute/Transpose算子
点击“阅读原文”，欢迎下载体验OneFlow新一代开源深度学习框架
