在使用pytorch訓(xùn)練模型，經(jīng)常需要加載大量圖片數(shù)據(jù)，因此pytorch提供了好用的數(shù)據(jù)加載工具Dataloader。

為了實現(xiàn)小批量循環(huán)讀取大型數(shù)據(jù)集，在Dataloader類具體實現(xiàn)中，使用了迭代器和生成器。

這一應(yīng)用場景正是python中迭代器模式的意義所在，因此本文對Dataloader中代碼進行解讀，可以更好的理解python中迭代器和生成器的概念。

本文的內(nèi)容主要有：

• 解釋python中的迭代器和生成器概念
• 解讀pytorch中Dataloader代碼，如何使用迭代器和生成器實現(xiàn)數(shù)據(jù)加載

python迭代基礎(chǔ)

python中圍繞著迭代有以下概念：

1. 可迭代對象 iterables
2. 迭代器 iterator
3. 生成器 generator

這三個概念互相關(guān)聯(lián)，并不是孤立的。在可迭代對象的基礎(chǔ)上發(fā)展了迭代器，在迭代器的基礎(chǔ)上又發(fā)展了生成器。

學(xué)習(xí)這些概念的名詞解釋沒有多大意義。編程中很多的抽象概念都是為了更好的實現(xiàn)某些功能，才去人為創(chuàng)造的協(xié)議和模式。

因此，要理解它們，需要探究概念背后的邏輯，為什么這樣設(shè)計？要解決的真正問題是什么？在哪些場景下應(yīng)用是最好的？

迭代模式首先要解決的基礎(chǔ)問題是，需要按一定順序獲取集合內(nèi)部數(shù)據(jù)，比如循環(huán)某個list。

當(dāng)數(shù)據(jù)很小時，不會有問題。但當(dāng)讀取大量數(shù)據(jù)時，一次性讀取會超出內(nèi)存限制，因此想出以下方法：

• 把大的數(shù)據(jù)分成幾個小塊，分批處理
• 惰性的取值方式，按需取值

循環(huán)讀數(shù)據(jù)可分為下面三種應(yīng)用場景，對應(yīng)著容器（可迭代對象），迭代器和生成器：

• for x in container: 為了遍歷python內(nèi)部序列容器（如list）, 這些類型內(nèi)部實現(xiàn)了__getitem__() 方法，可以從0開始按順序遍歷序列容器中的元素。
• for x in iterator: 為了循環(huán)用戶自定義的迭代器，需要實現(xiàn)__iter__和__next__方法，__iter__是迭代協(xié)議，具體每次迭代的執(zhí)行邏輯在 __next__或next方法里
• for x in generator: 為了節(jié)省循環(huán)的內(nèi)存和加速，使用生成器來實現(xiàn)惰性加載，在迭代器的基礎(chǔ)上加入了yield語句，最簡單的例子是 range(5)
  

代碼示例:

# 普通循環(huán) for x in list
numbers = [1, 2, 3,]
for n in numbers:
  print(n) # 1,2,3

# for循環(huán)實際干的事情
# iter輸入一個可迭代對象list，返回迭代器
# next方法取數(shù)據(jù)
my_iterator = iter(numbers)
next(my_iterator) # 1
next(my_iterator) # 2
next(my_iterator) # 3
next(my_iterator) # StopIteration exception

# 迭代器循環(huán) for x in iterator
for i,n in enumerate(numbers):
  print(i,n) # 0,1 / 1,3 / 2,3

# 生成器循環(huán) for x in generator
for i in range(3):
  print(i) # 0,1,2

上面示例代碼中python內(nèi)置函數(shù)iter和next的用法：

• iter函數(shù)，調(diào)用__iter__，返回一個迭代器
• next函數(shù)，輸入迭代器，調(diào)用__next__，取出數(shù)據(jù)

比較容易混淆的是__iter__和__next__兩個方法。它們的區(qū)別是：

• __iter__是為了可以迭代，真正執(zhí)行取數(shù)據(jù)的邏輯是__next__方法實現(xiàn)的，實際調(diào)用是通過next(iterator)完成
• __iter__可以返回自身（return self），實際讀取數(shù)據(jù)的實現(xiàn)放在__next__方法
• __iter__可以和yield搭配，返回生成器對象

__iter__返回自身的做法有點類似 python中的類型系統(tǒng)。為了保持一致性，python中一切皆對象。

每個對象創(chuàng)建后，都有類型指針，而類型對象的指針指向元對象，元對象的指針指向自身。

生成器，是在__iter__方法中加入yield語句，好處有：

• 減少循環(huán)判斷邏輯的復(fù)雜度
• 惰性取值，節(jié)省內(nèi)存和時間

yield作用：

• 代替函數(shù)中的return語句
• 記住上一次循環(huán)迭代器內(nèi)部元素的位置

三種循環(huán)模式常用函數(shù)

for x in container 方法:

• list, deque, …
• set, frozensets, …
• dict, defaultdict, OrderedDict, Counter, …
• tuple, namedtuple, …
• str

for x in iterator 方法:

• enumerate() # 加上list的index
• sorted() # 排序list
• reversed() # 倒序list
• zip() # 合并list

for x in generator 方法：

• range()
• map()
• filter()
• reduce()
• [x for x in list(...)]

Dataloder源碼分析

pytorch采用 for x in iterator 模式，從Dataloader類中讀取數(shù)據(jù)。

1. 為了實現(xiàn)該迭代模式，在Dataloader內(nèi)部實現(xiàn)__iter__方法，實際返回的是_DataLoaderIter類。
2. _DataLoaderIter類里面，實現(xiàn)了 __iter__方法，返回自身，具體執(zhí)行讀數(shù)據(jù)的邏輯，在__next__方法中。

以下代碼只截取了單線程下的數(shù)據(jù)讀取。

class DataLoader(object):
  r"""
  Data loader. Combines a dataset and a sampler, and provides
  single- or multi-process iterators over the dataset.
  """
  def __init__(self, dataset, batch_size=1, shuffle=False, ...):
    self.dataset = dataset
    self.batch_sampler = batch_sampler
    ...
  
  def __iter__(self):
    return _DataLoaderIter(self)

  def __len__(self):
    return len(self.batch_sampler)

class _DataLoaderIter(object):
  r"""Iterates once over the DataLoader's dataset, as specified by the sampler"""
  def __init__(self, loader):
    self.sample_iter = iter(self.batch_sampler)
    ...

  def __next__(self):
    if self.num_workers == 0: # same-process loading
      indices = next(self.sample_iter) # may raise StopIteration
      batch = self.collate_fn([self.dataset[i] for i in indices])
      if self.pin_memory:
        batch = pin_memory_batch(batch)
      return batch
    ...

  def __iter__(self):
    return self

Dataloader類中讀取數(shù)據(jù)Index的方法，采用了 for x in generator 方式，但是調(diào)用采用iter和next函數(shù)

1. 構(gòu)建隨機采樣類RandomSampler，內(nèi)部實現(xiàn)了 __iter__方法
2. __iter__方法內(nèi)部使用了 yield，循環(huán)遍歷數(shù)據(jù)集，當(dāng)數(shù)量達到batch_size大小時，就返回
3. 實例化隨機采樣類，傳入iter函數(shù)，返回一個迭代器
4. next會調(diào)用隨機采樣類中生成器，返回相應(yīng)的index數(shù)據(jù)

class RandomSampler(object):
  """random sampler to yield a mini-batch of indices."""
  def __init__(self, batch_size, dataset, drop_last=False):
    self.dataset = dataset
    self.batch_size = batch_size
    self.num_imgs = len(dataset)
    self.drop_last = drop_last

  def __iter__(self):
    indices = np.random.permutation(self.num_imgs)
    batch = []
    for i in indices:
      batch.append(i)
      if len(batch) == self.batch_size:
        yield batch
        batch = []
    ## if images not to yield a batch
    if len(batch)>0 and not self.drop_last:
      yield batch


  def __len__(self):
    if self.drop_last:
      return self.num_imgs // self.batch_size
    else:
      return (self.num_imgs + self.batch_size - 1) // self.batch_size

batch_sampler = RandomSampler(batch_size. dataset)
sample_iter = iter(batch_sampler)
indices = next(sample_iter)

總結(jié)

本文總結(jié)了python中循環(huán)的三種模式：

• for x in container 可迭代對象
• for x in iterator 迭代器
• for x in generator 生成器

pytorch中的數(shù)據(jù)加載模塊 Dataloader，使用生成器來返回數(shù)據(jù)的索引，使用迭代器來返回需要的張量數(shù)據(jù)，可以在大量數(shù)據(jù)情況下，實現(xiàn)小批量循環(huán)迭代式的讀取，避免了內(nèi)存不足問題。

參考文章

迭代器和生成器
流暢的Python-第14章：可迭代的對象、迭代器和生成器
pytorch-dataloader源碼

以上就是本文的全部內(nèi)容，希望對大家的學(xué)習(xí)有所幫助，也希望大家多多支持腳本之家。

最新評論