记录--深度学习加载大数据集遭遇内存和磁盘io瓶颈的问题

1--问题描述

       在最近的一个课题中，需要加载接近1T的数据集。由于内存的限制，无法一次加载整个数据集到内存中。

       想到的一个解决方案是，将大数据集处理为单个样本，根据DataLoader随机产生的索引index，通过一个list表来找到索引index对应的样本名及其路径path，最后当需要用到样本时，才会根据其路径path去加载数据。

        但上述的方法也会涉及内存瓶颈和磁盘io瓶颈的问题，当预加载到内存的数据不足时，程序就会处在等待的状态，等待从磁盘读取的数据传送至内存中，再从内存中读取数据。

        同时，由于博主存放数据集的服务器硬盘是机械硬盘，这就会导致从磁盘加载数据到内存的过程十分缓慢。而从内存读取的数据一旦传送到GPU上，就会被迅速处理掉，这也导致显卡长时间处于等待的空闲状态，GPU的利用率非常低，长时间处于0%。

        这时，博主想到的一个解决方法就是，减少数据在内存上的停留时间，即把CPU对数据集进行预处理的模块放在显卡上进行，不再由CPU放在内存中进行，而是直接传送到GPU上处理。通过这种方式缓解内存的压力，但终究治标不治本。

        由于内存的紧张，导致空闲的内存低，所以不能每次加载足够多的样本存放在内存中等待读取。同时，也会导致对磁盘的io请求十分频繁。

        这时候，也会产生一个疑问？即不能一次加载完整的数据集到内存中，而频繁的io请求来加载单个样本的数据，也会导致机械硬盘的io瓶颈。那为什么不把小部分的样本合并成一个大文件，一次就加载这个大文件放在内存中呢？博主也想过这个问题，但问题是博主读取训练样本的时候是随机产生索引的（即shuffle），我并不能保证一个batchsize中产生的样本索引都在这个大文件中，所以这种方法并不合适。

        在模型代码不能得到优化的前提下，目前想到的解决方法就是改用固态硬盘或者加大服务器的内存条容量。

        （写得比较乱。。。。。）

2--相关代码

①查看linux服务器内存情况

watch -n 1 free -h

②查看linux服务器磁盘io情况

iostat -x 1

③查看linux服务器显卡利用情况

watch -n 1 nvidia-smi

④查看用户进程情况

top -u username

top 1

3--相关解决方法

①参考给pytorch的Dataloader打鸡血：

参考链接1

参考链接2

②代码：

class data_prefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()
        self.preload()

    def preload(self):
        try:
            self.next_data, self.next_label, self.next_index = next(self.loader)
        except StopIteration:
            self.next_data = None
            self.next_label = None
            self.next_index = None
            return
        with torch.cuda.stream(self.stream):
            self.next_data = self.next_data.cuda(non_blocking=True)
            self.next_label = self.next_label.cuda(non_blocking=True)
            self.next_index = self.next_index.cuda(non_blocking=True)
            
    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        data = self.next_data
        label = self.next_label
        index = self.next_index
        self.preload()
        return data, label, index

######

self.data_loader['train'] = torch.utils.data.DataLoader(
    dataset=Feeder(**self.arg.train_feeder_args),
    batch_size=self.arg.batch_size,
    shuffle=True,
    num_workers=self.arg.num_worker,
    drop_last=True,
    worker_init_fn=init_seed,
    pin_memory=True)
loader = self.data_loader['train']
process = tqdm(loader, ncols=40)

#######

prefetcher = data_prefetcher(process)
rgb_data, label, index = prefetcher.next()
while rgb_data is not None:
    # forward
    ####
    # backward
    ####
    rgb_data, label, index = prefetcher.next()
          

经本人测试，这种prefetcher的方法无法解决io瓶颈；

4--参考

参考1

参考2

参考3

5--最新补充

博主将单个数据处理为单个小样本，并存放在固态硬盘（完美解决io瓶颈）中。根据Dataloader产生的索引去寻址加载数据，顺利解决上述问题，加载数据迅速，GPU利用率高。