简单记录一次dataloader低效运行DEBUG过程

key words: pytorch, dataloader, multiprocessing

问题描述

1. Dataset输入是一组对象，每个对象可以产生20-3000个训练样本。使用IterableDataset，用迭代的方式从每个对象中产生样本。为了支持多进程，预计算了每个对象会产生的样本数，使得分配给每个进程的对象产生的总样本数接近。
2. 一开始程序运行正常，开了num_workers=24，每个进程可以跑满100%的cpu利用率。
3. 然后发现了产生样本代码的bug，会重复访问图中的结点，添加visited数组阻止这种行为。
4. 重新运行程序，发现只有一个进程在满负荷运行，其他进程都处于等待状态。

解决记录

1. 排除了其他地方的代码造成区别的可能性，确定就是添加visited数组导致的。同时确定新增加的代码本身没有问题。
2. 猜测是DataLoader出了问题
3. 乱七八糟的各种尝试
4. 使用pycharm自带的profile工具，得到每个部分运行时间图，发现大部分时间都花在了DataLoader读取数据上，而其中的大部分时间都在进行等待。