Python迭代器和线程池使用示例

假如有一个列表，每次只取2个，不断迭代下去，直至取完，用Python迭代器实现。

def mynext(a, n):
    res = []

    for _ in range(n):
        try:
            res.append(next(a))
        except:
            break

    return res


if __name__ == '__main__':
    b = [1, 2, 3, 4, 5, 6, 7]
    c = iter(b)
    while True:
        d = mynext(c, 2)
        if len(d) == 0:
            break
        print(d)

运行结果：

业务需求：将迭代结果写入文件，比如深度学习中将PyTorch运行结果的热力图保存到文件中等。

import time


def write2file(str):
    print(str)
    time.sleep(2)


def mynext(a, n):
    res = []

    for _ in range(n):
        try:
            res.append(next(a))
        except:
            break

    return res


if __name__ == '__main__':
    b = [1, 2, 3, 4, 5, 6, 7]
    c = iter(b)

    begin = time.time()
    while True:
        d = mynext(c, 2)
        if len(d) == 0:
            break
        write2file(str(d))
    end = time.time()

    print(end - begin, 's')

运行结果：

注意：即使增大batch_size也不会缩减PyTorch测试时间，具体参考：为什么batch_size 增大会增加每一步的运算时间？ - 知乎，只能从写入文件这块入手。

用Python线程池加速上述代码。

import time
import multiprocessing


def write2file(str, cnt):
    print(cnt, str)
    time.sleep(2)
    print(cnt, 'OK')


def mynext(a, n):
    res = []

    for _ in range(n):
        try:
            res.append(next(a))
        except:
            break

    return res


if __name__ == '__main__':
    b = [1, 2, 3, 4, 5, 6, 7]
    c = iter(b)

    pool = multiprocessing.Pool(processes=4)

    begin = time.time()

    cnt = 0
    while True:
        d = mynext(c, 2)
        if len(d) == 0:
            break
        pool.apply_async(write2file, (str(d), cnt,))
        cnt += 1

    pool.close()
    pool.join()

    end = time.time()
    print(end - begin, 's')

运行结果：

参考：python进程池：multiprocessing.pool - jihite - 博客园、python线程池（threadpool）模块使用笔记 - Bypass - 博客园