完成第一版

Author: hhaAndroid

libv1/dataloader.py

@@ -41,12 +41,13 @@ def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,
             collate_fn = default_collate
         self.collate_fn = collate_fn
 
+    # 核心
     def __next__(self):
         index = next(self.batch_sampler)
         data = [self.dataset[idx] for idx in index]
-        data=self.collate_fn(data)
+        data = self.collate_fn(data)
         return data
 
+    # 返回自身，因为自身实现了 __next__
     def __iter__(self):
         return self
-

libv1/sampler.py

@@ -12,36 +12,14 @@ def __len__(self):
         raise NotImplementedError
 
 
-class BatchSampler(Sampler):
-    def __init__(self, sampler, batch_size, drop_last):
-        self.sampler = sampler
-        self.batch_size = batch_size
-        self.drop_last = drop_last
-
-    def __iter__(self):
-        batch = []
-        for idx in self.sampler:
-            batch.append(idx)
-            if len(batch) == self.batch_size:
-                yield batch
-                batch = []
-        if len(batch) > 0 and not self.drop_last:
-            yield batch
-
-    def __len__(self):
-        if self.drop_last:
-            return len(self.sampler) // self.batch_size
-        else:
-            return (len(self.sampler) + self.batch_size - 1) // self.batch_size
-
-
 class SequentialSampler(Sampler):
 
     def __init__(self, data_source):
         super(SequentialSampler, self).__init__(data_source)
         self.data_source = data_source
 
     def __iter__(self):
+        # 返回迭代器，不然无法 for .. in ..
         return iter(range(len(self.data_source)))
 
     def __len__(self):
@@ -51,8 +29,11 @@ def __len__(self):
 class RandomSampler(Sampler):
     def __init__(self, data_source, replacement=False, num_samples=None):
         super(RandomSampler, self).__init__(data_source)
+        # 数据集
         self.data_source = data_source
+        # 是否有放回抽象
         self.replacement = replacement
+        # 采样长度，一般等于 data_source 长度
         self._num_samples = num_samples
 
         if self._num_samples is not None and not replacement:
@@ -64,19 +45,50 @@ def __init__(self, data_source, replacement=False, num_samples=None):
                              "value, but got num_samples={}".format(self.num_samples))
 
     @property
-    def num_samples(self) -> int:
+    def num_samples(self):
         if self._num_samples is None:
             return len(self.data_source)
         return self._num_samples
 
     def __iter__(self):
         n = len(self.data_source)
+        # 通过 yield 关键字返回迭代器对象
         if self.replacement:
+            # 有放回抽样
+            # 可以直接写 yield from torch.randint(high=n, size=(self.num_samples,), dtype=torch.int64).tolist()
+            # 之所以按照每次生成32个，可能是因为想减少重复抽样概率 ?
             for _ in range(self.num_samples // 32):
                 yield from torch.randint(high=n, size=(32,), dtype=torch.int64).tolist()
             yield from torch.randint(high=n, size=(self.num_samples % 32,), dtype=torch.int64).tolist()
         else:
+            # 无放回抽样
             yield from torch.randperm(n).tolist()
 
     def __len__(self):
         return self.num_samples
+
+
+class BatchSampler(Sampler):
+    def __init__(self, sampler, batch_size, drop_last):
+        self.sampler = sampler
+        self.batch_size = batch_size
+        self.drop_last = drop_last
+
+    def __iter__(self):
+        batch = []
+        # 调用 sampler 内部的迭代器对象
+        for idx in self.sampler:
+            batch.append(idx)
+            # 如果已经得到了 batch 个 索引，则可以通过 yield 关键字生成生成器返回，得到迭代器对象
+            if len(batch) == self.batch_size:
+                yield batch
+                batch = []
+        if len(batch) > 0 and not self.drop_last:
+            yield batch
+
+    def __len__(self):
+        if self.drop_last:
+            # 如果最后的索引数不够一个 batch，则抛弃
+            return len(self.sampler) // self.batch_size
+        else:
+            return (len(self.sampler) + self.batch_size - 1) // self.batch_size

main.py

@@ -9,5 +9,64 @@ def demo_test_1():
         print(data)
 
 
+def demo_test_2():
+    from simplev1_datatset import SimpleV1Dataset
+    from libv1 import SequentialSampler, RandomSampler
+    from collections import Iterator, Iterable
+
+    simple_dataset = SimpleV1Dataset()
+    dataloader = DataLoader(simple_dataset, batch_size=2, collate_fn=default_collate)
+
+    print(isinstance(simple_dataset, Iterable))
+    print(isinstance(simple_dataset, Iterator))
+    print(isinstance(iter(simple_dataset), Iterator))
+
+    print(isinstance(SequentialSampler(simple_dataset), Iterable))
+    print(isinstance(SequentialSampler(simple_dataset), Iterator))
+    print(isinstance(iter(SequentialSampler(simple_dataset)), Iterator))
+
+    # BatchSampler 和 RandomSampler 内部实现结构一样，结果也是一样
+    print(isinstance(RandomSampler(simple_dataset), Iterable))
+    print(isinstance(RandomSampler(simple_dataset), Iterator))
+    print(isinstance(iter(RandomSampler(simple_dataset)), Iterator))
+
+    print(isinstance(dataloader, Iterator))
+
+
+def demo_test_3():
+    class DataLoader(object):
+        def __init__(self):
+            self.dataset = [[img0, target0], [img1, target1], [img2, target2], ..., [img99, target99]]
+            self.sampler = [0, 1, 2, 3, 4, ..., 99]
+            self.batch_size = 4
+            self.index = 0
+
+        def collate_fn(self, data):
+            batch_img = torch.stack(data[0], 0)
+            batch_target = torch.stack(data[1], 0)
+            return batch_img, batch_target
+
+        def __next__(self):
+            # 0.batch_index 输出
+            i = 0
+            batch_index = []
+            while i < self.batch_size:
+                batch_index.append(self.sampler[self.index])
+                self.index += 1
+                i += 1
+
+            # 1.得到 batch 个数据了
+            data = [self.dataset[idx] for idx in batch_index]
+
+            # 2.collate_fn 在 batch 维度拼接输出
+            batch_data = self.collate_fn(data)
+            return batch_data
+
+        def __iter__(self):
+            return self
+
+
 if __name__ == '__main__':
     demo_test_1()
+    # demo_test_2()
+