基于Python Yield的简易数据生成器为何运行速度异常缓慢？

嘿，我来帮你拆解下这个问题——从你给出的代码和性能测试结果来看，生成器运行缓慢的锅真不在Yield本身，核心问题出在你的数据加载逻辑和磁盘IO模式上，具体原因和优化方向如下：

核心慢因分析 #

• 单次磁盘IO的开销累积：你的loader函数每次调用都要单独读取一个pickle文件，单次耗时就有~40ms。如果你的生成器是逐个迭代调用loader（比如for mask in your_generator: ...），那每拿一个样本就触发一次磁盘读。当你要处理成百上千个样本时，几十毫秒的单次开销会迅速堆起来，整体速度自然慢得离谱。
• 小文件读取的低效性：50MB总数据拆成多个小pickle文件的话，单个文件肯定不大。操作系统对小文件的读取效率远低于大文件——每次读取都要走文件系统查找、磁盘寻道这些额外流程，这些开销占比甚至比实际读数据的时间还高，纯纯的浪费。
• Yield的惰性特性放大了延迟：Yield生成器是按需生成数据的，每调用一次next()才会生成一个样本。这就导致磁盘IO是串行、挨个触发的，完全没法利用磁盘的并行读取能力，也没机会预加载数据，直接把IO延迟拉满了。

针对性优化方案 #

方案1：批量合并pickle文件，减少IO次数 #

把多个样本打包到同一个pickle文件里，比如100个样本存一个文件，加载时一次读一批到内存，再通过生成器逐个返回：

def batch_loader(batch_size=100):
    # 获取所有pickle文件路径
    train_dir = os.path.join(CACHE_DIR, 'train')
    pickle_files = [f for f in os.listdir(train_dir) if f.endswith('.pickle')]
    
    for i in range(0, len(pickle_files), batch_size):
        batch_masks = []
        # 一次性读取一批文件
        for file in pickle_files[i:i+batch_size]:
            with open(os.path.join(train_dir, file), 'rb') as f:
                item = pickle.load(f)
                batch_masks.append(item['contour'])
        # 逐个返回批次内的样本
        for mask in batch_masks:
            yield mask

方案2：预加载全部数据到内存（最适合你的场景） #

你的总数据量只有50MB，完全可以一次性把所有pickle文件加载到内存里，之后生成器直接从内存取数据，彻底干掉磁盘IO开销：

def preloaded_generator(max_image_id):
    # 先把所有数据预加载到内存
    all_masks = []
    for image_id in range(max_image_id):
        cache_path = os.path.join(CACHE_DIR, 'train', '{:05}.pickle').format(image_id)
        with open(cache_path, 'rb') as f:
            item = pickle.load(f)
            all_masks.append(item['contour'])
    # 逐个返回内存中的数据
    for mask in all_masks:
        yield mask

用这个方法的话，生成器的单次迭代耗时会直接降到微秒级，速度提升几十倍甚至上百倍。

方案3：多线程预加载，隐藏IO延迟（适合大数据量场景） #

如果以后数据量变大到装不下内存，可以用多线程在后台预加载下一批数据，让磁盘IO和你的模型计算并行，把IO延迟藏起来：

import threading
from queue import Queue

def loader_worker(queue, image_ids):
    # 后台线程负责加载数据到队列
    for image_id in image_ids:
        cache_path = os.path.join(CACHE_DIR, 'train', '{:05}.pickle').format(image_id)
        with open(cache_path, 'rb') as f:
            item = pickle.load(f)
            queue.put(item['contour'])
    queue.put(None)  # 放个结束标记

def threaded_generator(image_ids, preload_size=10):
    queue = Queue(maxsize=preload_size)
    # 启动后台加载线程
    worker_thread = threading.Thread(target=loader_worker, args=(queue, image_ids))
    worker_thread.start()
    
    while True:
        mask = queue.get()
        if mask is None:
            break
        yield mask
    worker_thread.join()

优化效果验证 #

优化后你可以用%timeit测试下生成器的迭代速度，比如：

gen = preloaded_generator(max_image_id=201)  # 替换成你的实际样本数
%timeit next(gen)

此时单次迭代的耗时应该会从几十毫秒直接降到微秒级别，效果立竿见影。

内容的提问来源于stack exchange，提问作者YW P Kwon