| | |
| | import itertools |
| | from typing import Iterator, List, Optional, Sized, Union |
| |
|
| | import numpy as np |
| | import torch |
| | from mmengine.dataset import BaseDataset |
| | from mmengine.dist import get_dist_info, sync_random_seed |
| | from torch.utils.data import Sampler |
| |
|
| | from mmdet.registry import DATA_SAMPLERS |
| |
|
| |
|
| | @DATA_SAMPLERS.register_module() |
| | class MultiSourceSampler(Sampler): |
| | r"""Multi-Source Infinite Sampler. |
| | |
| | According to the sampling ratio, sample data from different |
| | datasets to form batches. |
| | |
| | Args: |
| | dataset (Sized): The dataset. |
| | batch_size (int): Size of mini-batch. |
| | source_ratio (list[int | float]): The sampling ratio of different |
| | source datasets in a mini-batch. |
| | shuffle (bool): Whether shuffle the dataset or not. Defaults to True. |
| | seed (int, optional): Random seed. If None, set a random seed. |
| | Defaults to None. |
| | |
| | Examples: |
| | >>> dataset_type = 'ConcatDataset' |
| | >>> sub_dataset_type = 'CocoDataset' |
| | >>> data_root = 'data/coco/' |
| | >>> sup_ann = '../coco_semi_annos/instances_train2017.1@10.json' |
| | >>> unsup_ann = '../coco_semi_annos/' \ |
| | >>> 'instances_train2017.1@10-unlabeled.json' |
| | >>> dataset = dict(type=dataset_type, |
| | >>> datasets=[ |
| | >>> dict( |
| | >>> type=sub_dataset_type, |
| | >>> data_root=data_root, |
| | >>> ann_file=sup_ann, |
| | >>> data_prefix=dict(img='train2017/'), |
| | >>> filter_cfg=dict(filter_empty_gt=True, min_size=32), |
| | >>> pipeline=sup_pipeline), |
| | >>> dict( |
| | >>> type=sub_dataset_type, |
| | >>> data_root=data_root, |
| | >>> ann_file=unsup_ann, |
| | >>> data_prefix=dict(img='train2017/'), |
| | >>> filter_cfg=dict(filter_empty_gt=True, min_size=32), |
| | >>> pipeline=unsup_pipeline), |
| | >>> ]) |
| | >>> train_dataloader = dict( |
| | >>> batch_size=5, |
| | >>> num_workers=5, |
| | >>> persistent_workers=True, |
| | >>> sampler=dict(type='MultiSourceSampler', |
| | >>> batch_size=5, source_ratio=[1, 4]), |
| | >>> batch_sampler=None, |
| | >>> dataset=dataset) |
| | """ |
| |
|
| | def __init__(self, |
| | dataset: Sized, |
| | batch_size: int, |
| | source_ratio: List[Union[int, float]], |
| | shuffle: bool = True, |
| | seed: Optional[int] = None) -> None: |
| |
|
| | assert hasattr(dataset, 'cumulative_sizes'),\ |
| | f'The dataset must be ConcatDataset, but get {dataset}' |
| | assert isinstance(batch_size, int) and batch_size > 0, \ |
| | 'batch_size must be a positive integer value, ' \ |
| | f'but got batch_size={batch_size}' |
| | assert isinstance(source_ratio, list), \ |
| | f'source_ratio must be a list, but got source_ratio={source_ratio}' |
| | assert len(source_ratio) == len(dataset.cumulative_sizes), \ |
| | 'The length of source_ratio must be equal to ' \ |
| | f'the number of datasets, but got source_ratio={source_ratio}' |
| |
|
| | rank, world_size = get_dist_info() |
| | self.rank = rank |
| | self.world_size = world_size |
| |
|
| | self.dataset = dataset |
| | self.cumulative_sizes = [0] + dataset.cumulative_sizes |
| | self.batch_size = batch_size |
| | self.source_ratio = source_ratio |
| |
|
| | self.num_per_source = [ |
| | int(batch_size * sr / sum(source_ratio)) for sr in source_ratio |
| | ] |
| | self.num_per_source[0] = batch_size - sum(self.num_per_source[1:]) |
| |
|
| | assert sum(self.num_per_source) == batch_size, \ |
| | 'The sum of num_per_source must be equal to ' \ |
| | f'batch_size, but get {self.num_per_source}' |
| |
|
| | self.seed = sync_random_seed() if seed is None else seed |
| | self.shuffle = shuffle |
| | self.source2inds = { |
| | source: self._indices_of_rank(len(ds)) |
| | for source, ds in enumerate(dataset.datasets) |
| | } |
| |
|
| | def _infinite_indices(self, sample_size: int) -> Iterator[int]: |
| | """Infinitely yield a sequence of indices.""" |
| | g = torch.Generator() |
| | g.manual_seed(self.seed) |
| | while True: |
| | if self.shuffle: |
| | yield from torch.randperm(sample_size, generator=g).tolist() |
| | else: |
| | yield from torch.arange(sample_size).tolist() |
| |
|
| | def _indices_of_rank(self, sample_size: int) -> Iterator[int]: |
| | """Slice the infinite indices by rank.""" |
| | yield from itertools.islice( |
| | self._infinite_indices(sample_size), self.rank, None, |
| | self.world_size) |
| |
|
| | def __iter__(self) -> Iterator[int]: |
| | batch_buffer = [] |
| | while True: |
| | for source, num in enumerate(self.num_per_source): |
| | batch_buffer_per_source = [] |
| | for idx in self.source2inds[source]: |
| | idx += self.cumulative_sizes[source] |
| | batch_buffer_per_source.append(idx) |
| | if len(batch_buffer_per_source) == num: |
| | batch_buffer += batch_buffer_per_source |
| | break |
| | yield from batch_buffer |
| | batch_buffer = [] |
| |
|
| | def __len__(self) -> int: |
| | return len(self.dataset) |
| |
|
| | def set_epoch(self, epoch: int) -> None: |
| | """Not supported in `epoch-based runner.""" |
| | pass |
| |
|
| |
|
| | @DATA_SAMPLERS.register_module() |
| | class GroupMultiSourceSampler(MultiSourceSampler): |
| | r"""Group Multi-Source Infinite Sampler. |
| | |
| | According to the sampling ratio, sample data from different |
| | datasets but the same group to form batches. |
| | |
| | Args: |
| | dataset (Sized): The dataset. |
| | batch_size (int): Size of mini-batch. |
| | source_ratio (list[int | float]): The sampling ratio of different |
| | source datasets in a mini-batch. |
| | shuffle (bool): Whether shuffle the dataset or not. Defaults to True. |
| | seed (int, optional): Random seed. If None, set a random seed. |
| | Defaults to None. |
| | """ |
| |
|
| | def __init__(self, |
| | dataset: BaseDataset, |
| | batch_size: int, |
| | source_ratio: List[Union[int, float]], |
| | shuffle: bool = True, |
| | seed: Optional[int] = None) -> None: |
| | super().__init__( |
| | dataset=dataset, |
| | batch_size=batch_size, |
| | source_ratio=source_ratio, |
| | shuffle=shuffle, |
| | seed=seed) |
| |
|
| | self._get_source_group_info() |
| | self.group_source2inds = [{ |
| | source: |
| | self._indices_of_rank(self.group2size_per_source[source][group]) |
| | for source in range(len(dataset.datasets)) |
| | } for group in range(len(self.group_ratio))] |
| |
|
| | def _get_source_group_info(self) -> None: |
| | self.group2size_per_source = [{0: 0, 1: 0}, {0: 0, 1: 0}] |
| | self.group2inds_per_source = [{0: [], 1: []}, {0: [], 1: []}] |
| | for source, dataset in enumerate(self.dataset.datasets): |
| | for idx in range(len(dataset)): |
| | data_info = dataset.get_data_info(idx) |
| | width, height = data_info['width'], data_info['height'] |
| | group = 0 if width < height else 1 |
| | self.group2size_per_source[source][group] += 1 |
| | self.group2inds_per_source[source][group].append(idx) |
| |
|
| | self.group_sizes = np.zeros(2, dtype=np.int64) |
| | for group2size in self.group2size_per_source: |
| | for group, size in group2size.items(): |
| | self.group_sizes[group] += size |
| | self.group_ratio = self.group_sizes / sum(self.group_sizes) |
| |
|
| | def __iter__(self) -> Iterator[int]: |
| | batch_buffer = [] |
| | while True: |
| | group = np.random.choice( |
| | list(range(len(self.group_ratio))), p=self.group_ratio) |
| | for source, num in enumerate(self.num_per_source): |
| | batch_buffer_per_source = [] |
| | for idx in self.group_source2inds[group][source]: |
| | idx = self.group2inds_per_source[source][group][ |
| | idx] + self.cumulative_sizes[source] |
| | batch_buffer_per_source.append(idx) |
| | if len(batch_buffer_per_source) == num: |
| | batch_buffer += batch_buffer_per_source |
| | break |
| | yield from batch_buffer |
| | batch_buffer = [] |
| |
|
