Spaces:
Running on Zero
Running on Zero
File size: 6,075 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """Data samplers."""
from __future__ import annotations
from collections.abc import Iterator
import numpy as np
from torch.utils.data import Dataset
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data.sampler import BatchSampler, Sampler
from mapdet3d.data.const import CommonKeys as K
from .datasets.base import VideoDataset
from .typing import DictDataOrList
class VideoInferenceSampler(
DistributedSampler[list[int]]
): # pragma: no cover # No unittest for distributed setting.
"""Produce sequence ordered indices for inference across all workers.
Inference needs to run on the __exact__ set of sequences and their
respective samples, therefore if the sequences are not divisible by the
number of workers or if they have different length, the sampler
produces different number of samples on different workers.
"""
def __init__(
self,
dataset: Dataset[DictDataOrList],
num_replicas: None | int = None,
rank: None | int = None,
shuffle: bool = True,
seed: int = 0,
drop_last: bool = False,
) -> None:
"""Creates an instance of the class.
Args:
dataset (Dataset): Inference dataset.
num_replicas (int, optional): Number of processes participating in
distributed training. By default, :attr:`world_size` is
retrieved from the current distributed group.
rank (int, optional): Rank of the current process within
:attr:`num_replicas`. By default, :attr:`rank` is retrieved
from the current distributed group.
shuffle (bool, optional): If ``True`` (default), sampler will
shuffle the indices.
seed (int, optional): random seed used to shuffle the sampler if
:attr:`shuffle=True`. This number should be identical across
all processes in the distributed group. Default: ``0``.
drop_last (bool, optional): if ``True``, then the sampler will drop
the tail of the data to make it evenly divisible across the
number of replicas. If ``False``, the sampler will add extra
indices to make the data evenly divisible across the replicas.
Default: ``False``.
"""
super().__init__(dataset, num_replicas, rank, shuffle, seed, drop_last)
assert hasattr(dataset, "video_mapping")
self.sequences = list(dataset.video_mapping["video_to_indices"])
self.num_seqs = len(self.sequences)
assert self.num_seqs >= self.num_replicas, (
f"Number of sequences ({self.num_seqs}) must be greater or "
f"equal to number of replicas ({self.num_replicas})!"
)
chunks = np.array_split(self.sequences, self.num_replicas)
self._local_seqs = chunks[self.rank]
self._local_idcs: list[int] = []
for seq in self._local_seqs:
self._local_idcs.extend(
dataset.video_mapping["video_to_indices"][seq]
)
def __iter__(self) -> Iterator[list[int]]:
"""Iteration method."""
return iter(self._local_idcs) # type: ignore
def __len__(self) -> int:
"""Return length of sampler instance."""
return len(self._local_idcs)
class AspectRatioBatchSampler(BatchSampler):
"""A sampler wrapper for grouping images with similar aspect ratio.
Moidified from:
https://github.com/open-mmlab/mmdetection/blob/main/mmdet/datasets/samplers/batch_sampler.py
Args:
sampler (Sampler): Base sampler.
batch_size (int): Size of mini-batch.
drop_last (bool): If ``True``, the sampler will drop the last batch if
its size would be less than ``batch_size``.
"""
def __init__(
self,
sampler: Sampler, # type: ignore
batch_size: int,
drop_last: bool = False,
) -> None:
"""Creates an instance of the class."""
if not isinstance(sampler, Sampler):
raise TypeError(
"sampler should be an instance of ``Sampler``, "
f"but got {sampler}"
)
super().__init__(sampler, batch_size, drop_last)
# two groups for w < h and w >= h
self._aspect_ratio_buckets: list[list[int]] = [[] for _ in range(2)]
def __iter__(self) -> Iterator[list[int]]:
"""Iteration method."""
for idx in self.sampler:
if hasattr(self.sampler, "dataset"):
data_dict = self.sampler.dataset[idx]
elif hasattr(self.sampler, "data_source"):
data_dict = self.sampler.data_source[idx]
else:
raise ValueError(
"sampler should have dataset or data_source attribute"
)
height, width = data_dict[K.input_hw]
bucket_id = 0 if width < height else 1
bucket = self._aspect_ratio_buckets[bucket_id]
bucket.append(idx)
# yield a batch of indices in the same aspect ratio group
if len(bucket) == self.batch_size:
yield bucket[:]
del bucket[:]
# yield the rest data and reset the bucket
left_data = (
self._aspect_ratio_buckets[0] + self._aspect_ratio_buckets[1]
)
self._aspect_ratio_buckets = [[] for _ in range(2)]
while len(left_data) > 0:
if len(left_data) <= self.batch_size:
if not self.drop_last:
yield left_data[:]
left_data = []
else:
yield left_data[: self.batch_size]
left_data = left_data[self.batch_size :]
def __len__(self) -> int:
"""Return length of sampler instance."""
if self.drop_last:
return len(self.sampler) // self.batch_size # type: ignore
return (
len(self.sampler) + self.batch_size - 1 # type: ignore
) // self.batch_size
|