完成第一版

Author: hhaAndroid

lib/__init__.py renamed to libv1/__init__.py

@@ -1,5 +1,6 @@
 from .dataset import Dataset
 from .sampler import Sampler, BatchSampler, SequentialSampler, RandomSampler
 from .dataloader import DataLoader
+from .collate import default_collate
 
-__all__ = ['Dataset', 'Sampler', 'BatchSampler', 'SequentialSampler', 'RandomSampler', 'DataLoader']
+__all__ = ['Dataset', 'Sampler', 'BatchSampler', 'SequentialSampler', 'RandomSampler', 'DataLoader', 'default_collate']

libv1/collate.py

@@ -0,0 +1,20 @@
+# -*- coding: utf-8 -*-
+# ======================================================
+# @Time    : 20-12-26 下午4:42
+# @Author  : huang ha
+# @Email   :
+# @File    : collate.py
+# @Comment: 
+# ======================================================
+import torch
+
+
+def default_collate(batch):
+    elem = batch[0]
+    elem_type = type(elem)
+    if isinstance(elem, torch.Tensor):
+        return torch.stack(batch, 0)
+    elif elem_type.__module__ == 'numpy':
+        return default_collate([torch.as_tensor(b) for b in batch])
+    else:
+        raise NotImplementedError

lib/dataloader.py renamed to libv1/dataloader.py

@@ -1,5 +1,4 @@
 from torch.utils.data._utils.collate import default_collate
-
 from .sampler import BatchSampler, SequentialSampler, RandomSampler
 
 
@@ -36,15 +35,18 @@ def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,
         self.batch_size = batch_size
         self.drop_last = drop_last
         self.sampler = sampler
-        self.batch_sampler = batch_sampler
+        self.batch_sampler = iter(batch_sampler)
 
         if collate_fn is None:
             collate_fn = default_collate
-
         self.collate_fn = collate_fn
 
-    @property
-    def _auto_collation(self):
-        return self.batch_sampler is not None
+    def __next__(self):
+        index = next(self.batch_sampler)
+        data = [self.dataset[idx] for idx in index]
+        data=self.collate_fn(data)
+        return data
 
+    def __iter__(self):
+        return self
 

lib/dataset.py renamed to libv1/dataset.py

@@ -13,8 +13,7 @@ def __len__(self):
 
 
 class BatchSampler(Sampler):
-    def __init__(self, sampler, batch_size, drop_last, data_source):
-        super(BatchSampler, self).__init__(data_source)
+    def __init__(self, sampler, batch_size, drop_last):
         self.sampler = sampler
         self.batch_size = batch_size
         self.drop_last = drop_last

lib/sampler.py renamed to libv1/sampler.py

@@ -1,16 +1,13 @@
-# This is a sample Python script.
+from libv1 import DataLoader, default_collate
 
-# Press Shift+F10 to execute it or replace it with your code.
-# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
 
-import torch.utils.data
-def print_hi(name):
-    # Use a breakpoint in the code line below to debug your script.
-    print(f'Hi, {name}')  # Press Ctrl+F8 to toggle the breakpoint.
+def demo_test_1():
+    from simplev1_datatset import SimpleV1Dataset
+    simple_dataset = SimpleV1Dataset()
+    dataloader = DataLoader(simple_dataset, batch_size=2, collate_fn=default_collate)
+    for data in dataloader:
+        print(data)
 
 
-# Press the green button in the gutter to run the script.
 if __name__ == '__main__':
-    print_hi('PyCharm')
-
-# See PyCharm help at https://www.jetbrains.com/help/pycharm/
+    demo_test_1()

main.py

@@ -0,0 +1,23 @@
+# -*- coding: utf-8 -*-
+# ======================================================
+# @Time    : 20-12-25 下午9:49
+# @Author  : huang ha
+# @Email   :
+# @File    : simple1_datatset.py
+# @Comment: 
+# ======================================================
+
+from libv1 import Dataset
+import numpy as np
+
+
+class SimpleV1Dataset(Dataset):
+    def __init__(self):
+        # 伪造数据
+        self.imgs = np.arange(0, 16).reshape(8, 2)
+
+    def __getitem__(self, index):
+        return self.imgs[index]
+
+    def __len__(self):
+        return self.imgs.shape[0]