帶你從零掌握迭代器及構(gòu)建最簡(jiǎn)DataLoader

from collections.abc import Iterable, Iteratorclass A(object):    def __init__(self):        self.a = [1, 2, 3]    def __iter__(self):        # 此處返回啥無(wú)所謂        return self.acls_a = A()#  Trueprint(isinstance(cls_a, Iterable))

for a in cls_a:    print(a)
# 程序報(bào)錯(cuò)，要理解這個(gè)錯(cuò)誤的含義TypeError: iter() returned non-iterator of type 'list'

class Iterable(metaclass=ABCMeta):
    # 如果實(shí)現(xiàn)了這個(gè)方法，那么就是 Iterable    @abstractmethod    def __iter__(self):        while False:            yield None
    @classmethod    def __subclasshook__(cls, C):        if cls is Iterable:            return _check_methods(C, "__iter__")        return NotImplemented

class Iterator(Iterable):
    # 迭代具體實(shí)現(xiàn)    @abstractmethod    def __next__(self):        'Return the next item from the iterator. When exhausted, raise StopIteration'        raise StopIteration
????# 返回自身，因?yàn)樽陨碛?__next__?方法(如果自身沒(méi)有?__next__，那么返回自身沒(méi)有意義)    def __iter__(self):????????return?self????????    @classmethod    def __subclasshook__(cls, C):        if cls is Iterator:            return _check_methods(C, '__iter__', '__next__')        return NotImplemented

class A(object):    def __init__(self):        self.index = -1        self.a = [1, 2, 3]
????#必須要返回一個(gè)實(shí)現(xiàn)了?__next__?方法的對(duì)象，否則后面無(wú)法?for?遍歷????#因?yàn)楸绢?lèi)自身實(shí)現(xiàn)了?__next__，所以通常都是返回?self?對(duì)象即可    def __iter__(self):        return self
    def __next__(self):        self.index += 1        if self.index < len(self.a):            return self.a[self.index]        else:????????????#拋異常，for?內(nèi)部會(huì)自動(dòng)捕獲，表示迭代完成            raise StopIteration("遍歷完了")cls_a = A()print(isinstance(cls_a, Iterable)) # Trueprint(isinstance(cls_a, Iterator)) # Trueprint(isinstance(iter(cls_a), Iterator)) # True
for a in cls_a:    print(a)# 打印 1 2 3

# 實(shí)際調(diào)用了?__iter__?方法返回自身，包括了?__next__?方法的對(duì)象cls_a = iter(cls_a)while True:    try:        # 然后調(diào)用對(duì)象的 __next__ 方法，不斷返回元素        value = next(cls_a)        print(value)    # 如果迭代完成，則捕獲異常即可    except StopIteration:        break

myiter = iter(cls_a)print(next(myiter))print(next(myiter))print(next(myiter))# 因?yàn)楸闅v完了，故此時(shí)會(huì)出現(xiàn)錯(cuò)誤：StopIteration: 遍歷完了print(next(myiter))

b=list([1,2,3])print(isinstance(b, Iterable)) # Trueprint(isinstance(b, Iterator)) # False

b=list([1,2,3])print(dir(b)) # 可以發(fā)現(xiàn)其存在 __iter__ 方法，不存在 __next__
b=iter(b) # 調(diào)用 list 內(nèi)部的 __iter__，返回了具備 __next__ 的對(duì)象print(isinstance(b, Iterable)) # Trueprint(isinstance(b, Iterator)) # Trueprint(dir(b)) # 同時(shí)具備 __iter__ 和 __next__ 方法

class A(object):    def __init__(self):        self.a = [1, 2, 3]    # 我們內(nèi)部又調(diào)用了 list 對(duì)象的 __iter__ 方法，故此時(shí)返回的對(duì)象是迭代器對(duì)象    def __iter__(self):        return iter(self.a)
cls_a = A()print(isinstance(cls_a, Iterable)) # Trueprint(isinstance(cls_a, Iterator)) # False
for a in cls_a:    print(a)# 輸出：1 2 3

# 僅僅實(shí)現(xiàn) __iter__ class A(object):    def __init__(self):        self.b = B()
    def __iter__(self):        return self.b
# 僅僅實(shí)現(xiàn) __next__class B(object):    def __init__(self):        self.index = -1        self.a = [1, 2, 3]
    def __next__(self):        self.index += 1        if self.index < len(self.a):            return self.a[self.index]        else:            # 內(nèi)部會(huì)自動(dòng)捕獲，表示迭代完成            raise StopIteration("遍歷完了")

cls_a = A()cls_b = B()print(isinstance(cls_a, Iterable)) # Trueprint(isinstance(cls_a, Iterator)) # Falseprint(isinstance(cls_b, Iterable)) # Falseprint(isinstance(cls_b, Iterator)) # False
print(type(iter(cls_a))) # B 對(duì)象print(isinstance(iter(cls_a), Iterator)) # False
for a in cls_a:    print(a)
# 輸出：1 2 3

class?A(object):    def __init__(self):        self.a = [1, 2, 3]
    def __getitem__(self, item):????????return?self.a[item]????????cls_a = A()print(isinstance(cls_a, Iterable))  # Falseprint(isinstance(cls_a, Iterator))  # Falseprint(dir(cls_a))  # 僅僅具備 __getitem__ 方法
cls_a = iter(cls_a)print(dir(cls_a))  # 具備 __iter__ 和 __next__ 方法
print(isinstance(cls_a, Iterable))  # Trueprint(isinstance(cls_a, Iterator))  # True
# 等價(jià)于?for?..?in?..while True:    try:????????# 然后調(diào)用對(duì)象的?__next__?方法，不斷返回元素        value = next(cls_a)        print(value)    # 如果迭代完成，則捕獲異常即可    except StopIteration:        break
# 輸出：1 2 3

a[0]?#?1??a[4] # 錯(cuò)誤，索引越界

class A(object):    def __init__(self):        self.a = [1, 2, 3]
    def __getitem__(self, item):        return self.a[item]
    def __len__(self):        return len(self.a)
cls_a = A()  print(len(cls_a)) # 3

def func():    for a in [1, 2, 3]:        yield a
cls_g = func()print(isinstance(cls_g, Iterator))  # Trueprint(dir(cls_g))  # 自動(dòng)具備 __iter__ 和 __next__ 方法
for a in cls_g:    print(a)
# 輸出: 1 2 3
# 一種更簡(jiǎn)單的寫(xiě)法是用 ()cls_g = (i for i in [1,2,3])

class DataLoader(object):    def __init__(self):        # 假設(shè)數(shù)據(jù)長(zhǎng)度是100，batch_size 是4        self.dataset = [[img0, target0], [img1, target1], ..., [img99, target99]]        # 假設(shè) sampler 是 SequentialSampler，那么實(shí)際上就是 [0,1,...,99] 列表而已        # 如果 sampler 是 RandomSampler，那么可能是 [30,1,34,2,6,...,0] 列表        self.sampler = [0, 1, 2, 3, 4, ..., 99]        self.batch_size = 4        self.index = 0
    def collate_fn(self, data):        # batch 維度聚合數(shù)據(jù)        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 輸出，實(shí)際上就是 BatchSampler 做的事情        i = 0        batch_index = []        while i < self.batch_size:            # 內(nèi)部會(huì)調(diào)用 sampler 對(duì)象取單個(gè)索引            batch_index.append(self.sampler[self.index])            self.index += 1            i += 1
        # 1.得到 batch 個(gè)數(shù)據(jù)了，調(diào)用 dataset 對(duì)象         data = [self.dataset[idx] for idx in batch_index]
        # 2. 調(diào)用 collate_fn 在 batch 維度拼接輸出        batch_data = self.collate_fn(data)        return batch_data
    def __iter__(self):        return self

from?simplev1_datatset?import?SimpleV1Datasetfrom 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)) # Falseprint(isinstance(simple_dataset, Iterator)) # Falseprint(isinstance(iter(simple_dataset), Iterator)) # True
print(isinstance(SequentialSampler(simple_dataset), Iterable)) # Trueprint(isinstance(SequentialSampler(simple_dataset), Iterator)) # Falseprint(isinstance(iter(SequentialSampler(simple_dataset)), Iterator)) # True
# BatchSampler?和?RandomSampler?內(nèi)部實(shí)現(xiàn)結(jié)構(gòu)一樣，結(jié)果也是一樣print(isinstance(RandomSampler(simple_dataset), Iterable)) # Trueprint(isinstance(RandomSampler(simple_dataset), Iterator)) # Falseprint(isinstance(iter(RandomSampler(simple_dataset)),?Iterator))?#?True
print(isinstance(dataloader, Iterator)) # True

def?__next__(self):????# 返回?batch?個(gè)索引????index?=?next(self.batch_sampler)????# 利用索引去取數(shù)據(jù)????data?=?[self.dataset[idx]?for?idx?in?index????# batch?維度聚合????data?=?self.collate_fn(data)????return?data

class?Dataset(object):????# 只要實(shí)現(xiàn)了?__getitem__?方法就可以變成迭代器    def __getitem__(self, index):          raise NotImplementedError     # 用于獲取數(shù)據(jù)集長(zhǎng)度     def __len__(self): ????????raise?NotImplementedError

class?Sampler(object):????def?__init__(self,?data_source):        pass  
    def __iter__(self):          raise NotImplementedError 
    def __len__(self):          raise NotImplementedError

class?SequentialSampler(Sampler): 
    def __init__(self, data_source):          super(SequentialSampler, self).__init__(data_source)         self.data_source = data_source 
    def __iter__(self):         # 返回迭代器，不然無(wú)法 for .. in ..         return iter(range(len(self.data_source))) 
    def __len__(self):         return len(self.data_source)

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 = []         # 調(diào)用 sampler 內(nèi)部的迭代器對(duì)象         for idx in self.sampler:             batch.append(idx)             # 如果已經(jīng)得到了 batch 個(gè) 索引，則可以通過(guò) yield             # 關(guān)鍵字生成生成器返回，得到迭代器對(duì)象             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:             # 如果最后的索引數(shù)不夠一個(gè) batch，則拋棄             return len(self.sampler) // self.batch_size ????????else:??????????????return? (len(self.sampler)?+?self.batch_size?-?1)?//?self.batch_size

class DataLoader(object):    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,?????????????????batch_sampler=None,?collate_fn=None,?drop_last=False)        self.dataset = dataset    
        # 因?yàn)檫@兩個(gè)功能是沖突的，假設(shè) shuffle=True,        # 但是 sampler 里面是 SequentialSampler，那么就違背設(shè)計(jì)思想了        if sampler is not None and shuffle:            raise ValueError('sampler option is mutually exclusive with '                             'shuffle')
        if batch_sampler is not None:            # 一旦設(shè)置了 batch_sampler，那么 batch_size、shuffle、sampler            # 和 drop_last 四個(gè)參數(shù)就不能傳入            # 因?yàn)檫@4個(gè)參數(shù)功能和 batch_sampler 功能沖突了            if batch_size != 1 or shuffle or sampler is not None or drop_last:                raise ValueError('batch_sampler option is mutually exclusive '                                 'with batch_size, shuffle, sampler, and '                                 'drop_last')            batch_size = None            drop_last = False
        if sampler is None:            if shuffle:                sampler = RandomSampler(dataset)            else:                sampler = SequentialSampler(dataset)
        # 也就是說(shuō) batch_sampler 必須要存在，你如果沒(méi)有設(shè)置，那么采用默認(rèn)類(lèi)        if batch_sampler is None:            batch_sampler = BatchSampler(sampler, batch_size, drop_last)
        self.batch_size = batch_size        self.drop_last = drop_last        self.sampler = sampler        self.batch_sampler = iter(batch_sampler)
        if collate_fn is 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)        return data
    # 返回自身，因?yàn)樽陨韺?shí)現(xiàn)了 __next__     def __iter__(self):        return self

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

class SimpleV1Dataset(Dataset):     def __init__(self):          # 偽造數(shù)據(jù)          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]  

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)