[Typing][A-97,A-99] Add type annotations for sampler and batch sampler (#66005)

SigureMo · web-flow · commit fb5f3eac6092 · 2024-07-13T22:21:30.000+08:00
diff --git a/python/paddle/io/dataloader/batch_sampler.py b/python/paddle/io/dataloader/batch_sampler.py
@@ -12,15 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import math
+from typing import (
+    Iterator,
+    Sequence,
+    Sized,
+)
 
 import numpy as np
 
 from .dataset import IterableDataset
 from .sampler import RandomSampler, Sampler, SequenceSampler
 
 
-class BatchSampler(Sampler):
+class BatchSampler(Sampler[Sequence[int]]):
     """
     A base implement of batch sampler used by `paddle.io.DataLoader`
     which yield mini-batch indices(a list/tuple with length as
@@ -98,14 +105,19 @@ class BatchSampler(Sampler):
             [53, 17, 22, 86, 52, 3, 92, 33]
     """
 
+    sampler: Sampler[int]
+    batch_size: int
+    shuffle: bool
+    drop_last: bool
+
     def __init__(
         self,
-        dataset=None,
-        sampler=None,
-        shuffle=False,
-        batch_size=1,
-        drop_last=False,
-    ):
+        dataset: Sized | None = None,
+        sampler: Sampler | None = None,
+        shuffle: bool = False,
+        batch_size: int = 1,
+        drop_last: bool = False,
+    ) -> None:
         if dataset is None:
             assert (
                 sampler is not None
@@ -142,7 +154,7 @@ def __init__(
         # in auto-parallel
         self._acc_steps = 1
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[list[int]]:
         local_batch_size = self.batch_size * self._acc_steps
         batch_indices = []
         for idx in self.sampler:
@@ -153,22 +165,25 @@ def __iter__(self):
         if not self.drop_last and len(batch_indices) > 0:
             yield batch_indices
 
-    def __len__(self):
+    def __len__(self) -> int:
         local_batch_size = self.batch_size * self._acc_steps
         num_samples = len(self.sampler)
         num_samples += int(not self.drop_last) * (local_batch_size - 1)
         return num_samples // local_batch_size
 
 
-class _InfiniteIterableSampler:
-    def __init__(self, dataset, batch_size=1):
+class _InfiniteIterableSampler(Sampler[Sequence[None]]):
+    dataset: IterableDataset
+    batch_size: int
+
+    def __init__(self, dataset: IterableDataset, batch_size: int = 1) -> None:
         assert isinstance(
             dataset, IterableDataset
         ), "dataset should be an instance of paddle.io.IterableDataset"
         self.dataset = dataset
         self.batch_size = batch_size
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[list[None]]:
         while True:
             yield [None] * self.batch_size
 
@@ -231,15 +246,24 @@ class DistributedBatchSampler(BatchSampler):
             ...     break
     """
 
+    dataset: Sized
+    batch_size: int
+    drop_last: bool
+    nranks: int
+    epoch: int
+    local_rank: int
+    num_samples: int
+    total_size: int
+
     def __init__(
         self,
-        dataset,
-        batch_size,
-        num_replicas=None,
-        rank=None,
-        shuffle=False,
-        drop_last=False,
-    ):
+        dataset: Sized,
+        batch_size: int,
+        num_replicas: int | None = None,
+        rank: int | None = None,
+        shuffle: bool = False,
+        drop_last: bool = False,
+    ) -> None:
         self.dataset = dataset
 
         assert (
@@ -279,7 +303,7 @@ def __init__(
         # in auto-parallel
         self._acc_steps = 1
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[list[int]]:
         local_batch_size = self.batch_size * self._acc_steps
         num_samples = len(self.dataset)
         indices = np.arange(num_samples).tolist()
@@ -337,13 +361,13 @@ def _get_indices_by_batch_size(indices):
         if not self.drop_last and len(batch_indices) > 0:
             yield batch_indices
 
-    def __len__(self):
+    def __len__(self) -> int:
         local_batch_size = self.batch_size * self._acc_steps
         num_samples = self.num_samples
         num_samples += int(not self.drop_last) * (local_batch_size - 1)
         return num_samples // local_batch_size
 
-    def set_epoch(self, epoch):
+    def set_epoch(self, epoch: int) -> None:
         """
         Sets the epoch number. When :attr:`shuffle=True`, this number is used
         as seeds of random numbers. By default, users may not set this, all
diff --git a/python/paddle/io/dataloader/sampler.py b/python/paddle/io/dataloader/sampler.py
@@ -12,13 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generator,
+    Generic,
+    Iterator,
+    Sequence,
+    Sized,
+    TypeVar,
+)
+
 import numpy as np
 
 from ...framework import core
 from ...tensor import randperm
 
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
+    from paddle import Tensor
+
+
+_T = TypeVar("_T")
 
-class Sampler:
+
+class Sampler(Generic[_T]):
     """
     An abstract class to encapsulate methods and behaviors of samplers.
 
@@ -85,17 +106,23 @@ class Sampler:
 
     """
 
-    def __init__(self, data_source=None):
+    data_source: Sized | None
+
+    def __init__(self, data_source: Sized | None = None) -> None:
         self.data_source = data_source
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[_T]:
         raise NotImplementedError
 
     # Not define __len__ method in this base class here for __len__
     # is not needed in same sence, e.g. paddle.io.IterableDataset
+    if TYPE_CHECKING:
+
+        def __len__(self) -> int:
+            ...
 
 
-class SequenceSampler(Sampler):
+class SequenceSampler(Sampler[int]):
     """
     Iterate samples sequentially, yield :code:`0, 1, 2, ..., len(data_source) -1`
     generally,
@@ -140,17 +167,19 @@ class SequenceSampler(Sampler):
     see `paddle.io.Sampler`
     """
 
-    def __init__(self, data_source):
+    data_source: Sized
+
+    def __init__(self, data_source: Sized) -> None:
         self.data_source = data_source
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[int]:
         return iter(range(len(self.data_source)))
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.data_source)
 
 
-class RandomSampler(Sampler):
+class RandomSampler(Sampler[int]):
     """
     Iterate samples randomly, yield shuffled indices, if :attr:`replacement=False`,
     yield shuffled indices of the whole data source, if :attr:`replacement=True`,
@@ -199,9 +228,17 @@ class RandomSampler(Sampler):
             87
     """
 
+    data_source: Sized
+    replacement: bool
+    generator: Generator[int, None, None] | None
+
     def __init__(
-        self, data_source, replacement=False, num_samples=None, generator=None
-    ):
+        self,
+        data_source: Sized,
+        replacement: bool = False,
+        num_samples: int | None = None,
+        generator: Generator[int, None, None] | None = None,
+    ) -> None:
         self.data_source = data_source
         self.replacement = replacement
         self._num_samples = num_samples
@@ -226,12 +263,12 @@ def __init__(
             )
 
     @property
-    def num_samples(self):
+    def num_samples(self) -> int:
         if self._num_samples is None:
             return len(self.data_source)
         return self._num_samples
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[int]:
         n = len(self.data_source)
         if self.generator:
             for i in range(self.num_samples):
@@ -252,7 +289,7 @@ def __iter__(self):
                 ).tolist():
                     yield index
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.num_samples
 
 
@@ -288,7 +325,7 @@ def _weighted_sample(weights, num_samples, replacement=True):
     return np.array(rets)
 
 
-class WeightedRandomSampler(Sampler):
+class WeightedRandomSampler(Sampler[int]):
     """
     Random sample with given weights (probabilities), sample index will be in range
     [0, len(weights) - 1], if :attr:`replacement` is True, index can be sampled
@@ -325,7 +362,16 @@ class WeightedRandomSampler(Sampler):
             1
     """
 
-    def __init__(self, weights, num_samples, replacement=True):
+    weights: npt.NDArray[Any] | Tensor | Sequence[float]
+    num_samples: int
+    replacement: bool
+
+    def __init__(
+        self,
+        weights: npt.NDArray[Any] | Tensor | Sequence[float],
+        num_samples: int,
+        replacement: bool = True,
+    ) -> None:
         if not isinstance(num_samples, int) or num_samples <= 0:
             raise ValueError("num_samples should be a positive integer")
         if not isinstance(replacement, bool):
@@ -334,18 +380,18 @@ def __init__(self, weights, num_samples, replacement=True):
         self.num_samples = num_samples
         self.replacement = replacement
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[int]:
         idxs = _weighted_sample(
             self.weights, self.num_samples, self.replacement
         )
         return iter(idxs.reshape(-1).tolist())
 
-    def __len__(self):
+    def __len__(self) -> int:
         mul = np.prod(self.weights.shape) // self.weights.shape[-1]
         return self.num_samples * mul
 
 
-class SubsetRandomSampler(Sampler):
+class SubsetRandomSampler(Sampler[int]):
     r"""
     Randomly sample elements from a given list of indices, without replacement.
 
@@ -372,14 +418,16 @@ class SubsetRandomSampler(Sampler):
 
     """
 
-    def __init__(self, indices):
+    indices: Sequence[int]
+
+    def __init__(self, indices: Sequence[int]) -> None:
         if len(indices) == 0:
             raise ValueError(
                 "The length of `indices` in SubsetRandomSampler should be greater than 0."
             )
         self.indices = indices
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[int]:
         for i in randperm(len(self.indices)):
             yield self.indices[i]
 
