feature: add sequentialsampler

Author: FelixAbrahamsson

datastream/datastream.py

@@ -505,3 +505,16 @@ def test_take():
         Datastream(Dataset.from_subscriptable(list('d'))),
     ])
     assert len(list(datastream.take(2).data_loader(batch_size=1))) == 2
+
+
+def test_sequential_sampler():
+
+    from datastream.samplers import SequentialSampler
+
+    dataset = Dataset.from_subscriptable(list('abc'))
+    datastream = Datastream(dataset, SequentialSampler(len(dataset))).take(2)
+    assert len(list(datastream.data_loader(batch_size=1))) == 2
+
+    datastream = Datastream(dataset, SequentialSampler(len(dataset)))
+    it = iter(datastream.data_loader(batch_size=6, n_batches_per_epoch=10))
+    assert next(it) == ['a', 'b', 'c', 'a', 'b', 'c']

datastream/samplers/__init__.py

@@ -1,4 +1,5 @@
 from datastream.samplers.standard_sampler import StandardSampler
+from datastream.samplers.sequential_sampler import SequentialSampler
 from datastream.samplers.merge_sampler import MergeSampler
 from datastream.samplers.multi_sampler import MultiSampler
 from datastream.samplers.repeat_sampler import RepeatSampler

datastream/samplers/sequential_sampler.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+from pydantic import BaseModel
+import torch
+
+
+class SequentialSampler(BaseModel, torch.utils.data.Sampler):
+    sampler: torch.utils.data.SequentialSampler
+
+    class Config:
+        arbitrary_types_allowed = True
+        allow_mutation = False
+
+    def __init__(self, length):
+        BaseModel.__init__(
+            self,
+            sampler=torch.utils.data.SequentialSampler(torch.ones(length))
+        )
+
+    def __len__(self):
+        return len(self.sampler)
+
+    def __iter__(self):
+        return iter(self.sampler)
+
+    def sample_proportion(self, proportion):
+        sampler = SequentialSampler(int(len(self) * proportion))
+        return sampler