paddle.io.BatchSampler support iterable sampler 易用性提升 (#66170)

NKNaN · web-flow · commit 42a4bafb41b9 · 2024-07-19T18:49:15.000+08:00
* update batchsampler

* update test
diff --git a/python/paddle/io/dataloader/batch_sampler.py b/python/paddle/io/dataloader/batch_sampler.py
@@ -16,6 +16,7 @@
 
 import math
 from typing import (
+    Iterable,
     Iterator,
     Sequence,
     Sized,
@@ -47,7 +48,7 @@ class BatchSampler(Sampler[Sequence[int]]):
                 :ref:`api_paddle_io_IterableDataset` or other python object which implemented
                 :code:`__len__` for BatchSampler to get indices as the
                 range of :attr:`dataset` length. Default None, disabled.
-        sampler (Sampler, optional): this should be a :ref:`api_paddle_io_Sample`
+        sampler (Sampler, Iterable, optional): this should be a :ref:`api_paddle_io_Sample` or Iterable
                 instance which implemented :code:`__iter__` to generate
                 sample indices. :attr:`sampler` and :attr:`dataset`
                 can not be set in the same time.  If :attr:`sampler`
@@ -105,15 +106,15 @@ class BatchSampler(Sampler[Sequence[int]]):
             [53, 17, 22, 86, 52, 3, 92, 33]
     """
 
-    sampler: Sampler[int]
+    sampler: Sampler[int] | Iterable[int]
     batch_size: int
     shuffle: bool
     drop_last: bool
 
     def __init__(
         self,
         dataset: Sized | None = None,
-        sampler: Sampler | None = None,
+        sampler: Sampler | Iterable[int] | None = None,
         shuffle: bool = False,
         batch_size: int = 1,
         drop_last: bool = False,
@@ -123,8 +124,8 @@ def __init__(
                 sampler is not None
             ), "either dataset or sampler should be set"
             assert isinstance(
-                sampler, Sampler
-            ), f"sampler should be a paddle.io.Sampler, but got {type(sampler)}"
+                sampler, (Sampler, Iterable)
+            ), f"sampler should be either paddle.io.Sampler or Iterable, but got {type(sampler)}"
             assert not shuffle, "shuffle should be False when sampler is set"
             self.sampler = sampler
         else:
diff --git a/test/legacy_test/test_batch_sampler.py b/test/legacy_test/test_batch_sampler.py
@@ -245,6 +245,39 @@ def test_main(self):
             pass
 
 
+class TestBatchSamplerWithIterableSampler(TestBatchSampler):
+    def init_batch_sampler(self):
+        sampler = range(1000)
+        bs = BatchSampler(
+            sampler=sampler,
+            batch_size=self.batch_size,
+            drop_last=self.drop_last,
+        )
+        return bs
+
+
+class TestBatchSamplerWithIterableSamplerDropLast(
+    TestBatchSamplerWithIterableSampler
+):
+    def setUp(self):
+        self.num_samples = 1000
+        self.num_classes = 10
+        self.batch_size = 32
+        self.shuffle = False
+        self.drop_last = True
+
+
+class TestBatchSamplerWithIterableSamplerShuffle(
+    TestBatchSamplerWithIterableSampler
+):
+    def setUp(self):
+        self.num_samples = 1000
+        self.num_classes = 10
+        self.batch_size = 32
+        self.shuffle = True
+        self.drop_last = True
+
+
 class TestWeightedRandomSampler(unittest.TestCase):
     def init_probs(self, total, pos):
         pos_probs = np.random.random((pos,)).astype('float32')
