# Copyright (c) Alibaba, Inc. and its affiliates.
import platform
import random
from distutils.version import LooseVersion
from functools import partial
import numpy as np
import torch
from mmcv.parallel import collate
from mmcv.runner import get_dist_info
from torch.utils.data import DataLoader
from easycv.datasets.builder import build_sampler
from easycv.datasets.shared.odps_reader import set_dataloader_workid
from easycv.framework.errors import NotImplementedError
from easycv.utils.dist_utils import sync_random_seed
from easycv.utils.torchacc_util import is_torchacc_enabled
from .collate import CollateWrapper
if platform.system() != 'Windows':
# https://github.com/pytorch/pytorch/issues/973
import resource
rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1]))
[docs]def build_dataloader(dataset,
imgs_per_gpu,
workers_per_gpu,
num_gpus=1,
dist=True,
shuffle=True,
replace=False,
seed=None,
reuse_worker_cache=False,
odps_config=None,
persistent_workers=False,
collate_hooks=None,
use_repeated_augment_sampler=False,
sampler=None,
pin_memory=False,
**kwargs):
"""Build PyTorch DataLoader.
In distributed training, each GPU/process has a dataloader.
In non-distributed training, there is only one dataloader for all GPUs.
Args:
dataset (Dataset): A PyTorch dataset.
imgs_per_gpu (int): Number of images on each GPU, i.e., batch size of
each GPU.
workers_per_gpu (int): How many subprocesses to use for data loading
for each GPU.
num_gpus (int): Number of GPUs. Only used in non-distributed training.
dist (bool): Distributed training/test or not. Default: True.
shuffle (bool): Whether to shuffle the data at every epoch.
Default: True.
replace (bool): Replace or not in random shuffle.
It works on when shuffle is True.
seed (int, Optional): The seed. Default to None.
reuse_worker_cache (bool): If set true, will reuse worker process so that cached
data in worker process can be reused.
persistent_workers (bool) : After pytorch1.7, could use persistent_workers=True to
avoid reconstruct dataworker before each epoch, speed up before epoch
use_repeated_augment_sampler (bool) : If set true, it will use RASampler.
Default: False.
kwargs: any keyword argument to be used to initialize DataLoader
Returns:
DataLoader: A PyTorch dataloader.
"""
rank, world_size = get_dist_info()
if dist:
seed = sync_random_seed(seed)
batch_size = imgs_per_gpu
num_workers = workers_per_gpu
else:
if replace:
raise NotImplementedError
batch_size = num_gpus * imgs_per_gpu
num_workers = num_gpus * workers_per_gpu
default_sampler_args = dict(
dataset=dataset,
num_replicas=world_size,
rank=rank,
shuffle=shuffle,
seed=seed,
replace=replace)
split_huge_listfile_byrank = getattr(dataset, 'split_huge_listfile_byrank',
False)
if sampler is not None:
sampler_cfg = sampler
sampler_cfg.update(default_sampler_args)
elif use_repeated_augment_sampler:
sampler_cfg = dict(type='RASampler', **default_sampler_args)
elif hasattr(dataset, 'm_per_class') and dataset.m_per_class > 1:
sampler_cfg = dict(
type='DistributedMPSampler',
split_huge_listfile_byrank=split_huge_listfile_byrank,
**default_sampler_args)
else:
if dist:
sampler_cfg = dict(
type='DistributedSampler',
split_huge_listfile_byrank=split_huge_listfile_byrank,
**default_sampler_args)
else:
sampler_cfg = dict(
type='RandomSampler',
data_source=dataset) if shuffle else None # TODO: set replace
sampler = build_sampler(sampler_cfg) if sampler_cfg is not None else None
init_fn = partial(
worker_init_fn,
num_workers=num_workers,
rank=rank,
seed=seed,
odps_config=odps_config) if seed is not None else None
collate_fn = dataset.collate_fn if hasattr(
dataset, 'collate_fn') else partial(
collate, samples_per_gpu=imgs_per_gpu)
if collate_hooks:
collate_fn = CollateWrapper(collate_fn, collate_hooks)
if not reuse_worker_cache:
if LooseVersion(torch.__version__) < LooseVersion('1.7.0'):
print(
'Pytorch Version < 1.7, build Dataloader without persistent_workers'
)
data_loader = DataLoader(
dataset,
batch_size=batch_size,
sampler=sampler,
num_workers=num_workers,
collate_fn=collate_fn,
pin_memory=pin_memory,
worker_init_fn=init_fn,
**kwargs)
else:
data_loader = DataLoader(
dataset,
batch_size=batch_size,
sampler=sampler,
num_workers=num_workers,
collate_fn=collate_fn,
pin_memory=pin_memory,
worker_init_fn=init_fn,
persistent_workers=persistent_workers,
**kwargs)
else:
# use InfiniteDataLoader to reuse worker process for caching data
data_loader = InfiniteDataLoader(
dataset,
batch_size=batch_size,
sampler=sampler,
num_workers=num_workers,
collate_fn=collate_fn,
pin_memory=pin_memory,
worker_init_fn=init_fn,
**kwargs)
if is_torchacc_enabled():
from .loader_wrapper import TorchaccLoaderWrapper
data_loader = TorchaccLoaderWrapper(data_loader)
return data_loader
[docs]def worker_init_fn(worker_id, num_workers, rank, seed, odps_config=None):
# The seed of each worker equals to
# num_worker * rank + worker_id + user_seed
worker_seed = num_workers * rank + worker_id + seed
np.random.seed(worker_seed)
random.seed(worker_seed)
torch.manual_seed(worker_seed)
if odps_config is not None:
# for odps to set correct offset in multi-process pytorch dataloader
# use init_fn to set global DATALOADER_WORKID before dataset getitem
set_dataloader_workid(worker_id)
# set_dataloader_worknum(imgs_per_gpu)
[docs]class InfiniteDataLoader(torch.utils.data.dataloader.DataLoader):
""" Dataloader that reuses workers. https://github.com/pytorch/pytorch/issues/15849
Uses same syntax as vanilla DataLoader.
"""
[docs] def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
object.__setattr__(self, 'batch_sampler',
_RepeatSampler(self.batch_sampler))
self.iterator = super().__iter__()
def __len__(self):
return len(self.batch_sampler.sampler)
def __iter__(self):
for i in range(len(self)):
yield next(self.iterator)
class _RepeatSampler(object):
""" Sampler that repeats forever.
Args:
sampler (Sampler)
"""
def __init__(self, sampler):
self.sampler = sampler
def __iter__(self):
while True:
yield from iter(self.sampler)