Python-迭代器与生成器

迭代器

迭代是Python最强大的功能之一，是访问集合元素如list，tuple的一种方式。迭代器对象从集合的第一个元素开始访问，直到所有的元素被访问完结束。迭代器只能往前不会后退。
迭代器有两个基本的方法：iter() 和 next()。
迭代器最常见的就是引用于深度学习中，一个batch一个batch的读数据，避免内存爆掉的情况，比如Pytorch中最常用的Dataloader(差点忘了经常自定义的dataloader是迭代器的原理)
如：

迭代器对象可以使用常规for语句进行遍历,不需要担心越界问题，因为迭代器内部存在着迭代终止标记StopIteration。

迭代器有什么好处？

在for循环中，如果是python2，则需要系统一次性分配相应长度的空间给list，而在python3中，只需要一个range对象，所需值由这个对象内置的next方法迭代返回，节省了空间。

创建一个迭代器

想要自定义一个迭代器需要在你的类中实现两个方法:iter() 与 next() 。

• iter() 方法返回一个特殊的迭代器对象， 这个迭代器对象实现了 next() 方法并通过 StopIteration 异常标识迭代的完成。
• next() 方法会返回迭代器的输出。

class MyNumbers:
    def __iter__(self):
        self.a = 1
        return self

    def __next__(self):
        x = self.a
        self.a += 1
        return x

myclass = MyNumbers()
myiter = iter(myclass)

print(next(myiter))
print(next(myiter))
print(next(myiter))
print(next(myiter))
print(next(myiter))

1
2
3
4
5

当然，这样一个迭代器是危险的，因为其没有判断什么时候迭代停止，如果运用在for循环中会导致无限循环。这时候就需要一个StopIteration
StopIteration 用于标识迭代的完成，防止出现无限循环的情况，在 next() 方法中我们可以设置在完成指定循环次数后触发 StopIteration 异常来结束迭代。

class MyNumbers:
  def __iter__(self):
    self.a = 1
    return self

  def __next__(self):
    if self.a <= 20:
      x = self.a
      self.a += 1
      return x
    else:
      raise StopIteration

myclass = MyNumbers()
myiter = iter(myclass)

生成器

在 Python 中，使用了 yield 的函数被称为生成器（generator）。
跟普通函数不同的是，生成器是一个返回迭代器的函数，只能用于迭代操作，更简单点理解生成器就是一个迭代器。
在调用生成器运行的过程中，每次遇到 yield 时函数会暂停并保存当前所有的运行信息，返回 yield 的值, 并在下一次执行 next() 方法时从当前位置继续运行。
调用一个生成器函数，返回的是一个迭代器对象。

import sys
def fibonacci(n): # 生成器函数 - 斐波那契
    a, b, counter = 0, 1, 0
    while True:
        if (counter > n): 
            return
        yield a
        a, b = b, a + b
        counter += 1
f = fibonacci(10) # f 是一个迭代器，由生成器返回生成
while True:
    try:
        print (next(f), end=" ")
    except StopIteration:
        sys.exit()

Pytorch中的生成器的应用

自定义数据加载器的时候，我们需要继承Dataset类并重写__len__和__getitem__ 这两个方法

import torch
class Dataset(torch.utils.data.Dataset):
  'Characterizes a dataset for PyTorch'
  def __init__(self, list_IDs, labels):
        'Initialization'
        self.labels = labels
        self.list_IDs = list_IDs

  def __len__(self):
        'Denotes the total number of samples'
        return len(self.list_IDs)

  def __getitem__(self, index):
        'Generates one sample of data'
        # Select sample
        ID = self.list_IDs[index]

        # Load data and get label
        X = torch.load('data/' + ID + '.pt')
        y = self.labels[ID]

        return X, y
#生成数据generator
training_set = Dataset(partition['train'], labels)
training_generator = torch.utils.data.DataLoader(training_set, **params)

通过将整个Dataset传递给Dataloader，Dataloader内部会有yield操作来返回数据


具体代码可以看
https://github.com/pytorch/pytorch/blob/master/torch/utils/data/dataloader.py