# 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, RandomSampler
from easycv.datasets.shared.odps_reader import set_dataloader_workid
from .sampler import DistributedMPSampler, DistributedSampler
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,
**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.
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
kwargs: any keyword argument to be used to initialize DataLoader
Returns:
DataLoader: A PyTorch dataloader.
"""
if dist:
rank, world_size = get_dist_info()
split_huge_listfile_byrank = getattr(dataset,
'split_huge_listfile_byrank',
False)
if hasattr(dataset, 'm_per_class') and dataset.m_per_class > 1:
sampler = DistributedMPSampler(
dataset,
world_size,
rank,
shuffle=shuffle,
split_huge_listfile_byrank=split_huge_listfile_byrank)
else:
sampler = DistributedSampler(
dataset,
world_size,
rank,
shuffle=shuffle,
split_huge_listfile_byrank=split_huge_listfile_byrank)
batch_size = imgs_per_gpu
num_workers = workers_per_gpu
else:
if replace:
raise NotImplementedError
if hasattr(dataset, 'm_per_class') and dataset.m_per_class > 1:
sampler = DistributedMPSampler(
dataset, 1, 0, shuffle=shuffle, replace=replace)
else:
sampler = RandomSampler(
dataset) if shuffle else None # TODO: set replace
batch_size = num_gpus * imgs_per_gpu
num_workers = num_gpus * workers_per_gpu
init_fn = partial(worker_init_fn, seed=seed, odps_config=odps_config)
collate_fn = dataset.collate_fn if hasattr(
dataset, 'collate_fn') else partial(
collate, samples_per_gpu=imgs_per_gpu)
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=False,
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=False,
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=False,
worker_init_fn=init_fn,
**kwargs)
return data_loader
[docs]def worker_init_fn(worker_id, seed=None, odps_config=None):
if seed is not None:
worker_seed = worker_id + seed
np.random.seed(worker_seed)
random.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)