epoch训练时间不同_给训练踩踩油门——Pytorch加速数据读取

加入极市专业CV交流群，与6000+来自腾讯，华为，百度，北大，清华，中科院等名企名校视觉开发者互动交流！更有机会与李开复老师等大牛群内互动！

同时提供每月大咖直播分享、真实项目需求对接、干货资讯汇总，行业技术交流。关注 极市平台 公众号 ，回复 加群，立刻申请入群~

本文选自知乎，已获得作者授权，不得二次转载。

作者：MrTian

链接：https://zhuanlan.zhihu.com/p/80695364

需求

最近在训练coco数据集，训练集就有11万张，训练一个epoch就要将近100分钟，训练100个epoch，就需要7天！这实在是太慢了。

经过观察，发现训练时GPU利用率不是很稳定，每训练5秒，利用率都要从100%掉到0%一两秒，初步判断是数据读取那块出现了瓶颈。于是经过调研和实验，制定了下列解决方案。

解决方案

安装：

pip install prefetch_generator

使用：

from torch.utils.data import DataLoader

from prefetch_generator import BackgroundGenerator

class DataLoaderX(DataLoader):

    def __iter__(self):

        return BackgroundGenerator(super().__iter__())

然后用DataLoaderX替换原本的DataLoader。

提速原因：

原本Pytorch默认的DataLoader会创建一些worker线程来预读取新的数据，但是除非这些线程的数据全部都被清空，这些线程才会读下一批数据。使用prefetch_generator，我们可以保证线程不会等待，每个线程都总有至少一个数据在加载。

(2)data_prefetcher

使用data_prefetcher新开cuda stream来拷贝tensor到gpu。

使用：

    def __init__(self, loader, opt):

        self.loader = iter(loader)