pytorch数据加载程序非常慢的第一个历元

Question

当我创建一个PyTorch数据加载器并开始迭代时——我得到一个非常慢的第一个历元（x10--x30比所有下一个历元都慢）。此外，只有Kaggle的Google landmark recognition 2020中的火车数据集才会出现此问题。我无法在合成图像上重现这一点，同时，我尝试创建一个包含来自GLR2020的500k图像的文件夹，并且一切都很顺利。在PyTorch论坛中发现了一些类似的问题，但没有任何解决方案

import argparse
import pandas as pd
import numpy as np
import os, sys
import multiprocessing, ray
import time
import cv2
import logging
import albumentations as albu
from torch.utils.data import Dataset, DataLoader

samples = 50000 # count of samples to speed up test
bs = 64 # batch size
dir = '/hdd0/datasets/ggl_landmark_recognition_2020/train' # directory with train data
all_files = pd.read_csv('/hdd0/datasets/ggl_landmark_recognition_2020/train.csv')
files = np.random.choice(all_files.id.values, 50000)
files = [os.path.join(_[0], _[1], _[2], _+'.jpg') for _ in files]

# augmentations
aug =  albu.Compose([albu.Resize(400, 400),
        albu.Rotate(limit=15),
        albu.ChannelDropout(p=0.1),
        albu.Normalize(),])

class ImgDataset:
    def __init__(self, path, files, augmentation = None):
        self.path = path
        self.files = {k:v for k, v in enumerate(files)}
        self.augmentation = augmentation

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        img_name = self.files[idx]
        img = np.array(cv2.imread(os.path.join(self.path, img_name)))
        if self.augmentation is not None:
            return self.augmentation(image=img)['image']


dtset = ImgDataset(dir,files, aug)
torchloader = DataLoader(dataset= dtset, batch_size=64, num_worker=16, shuffle=True)
for _ in range(3):
   t1 = time.time()
   for idx, val in enumerate(torchloader):
       pass
   t2 = time.time()
   print(str(t2-t1) +' sec')

下面是一些在DataLoader中使用不同num_workers的执行速度示例

#num_workers=0
273.1584792137146 sec
83.15653467178345 sec
83.67923021316528 sec

# num_workers = 8 
165.62366938591003 sec
10.405716896057129 sec
10.495309114456177 sec

# num_workers = 16
156.60744667053223 sec
8.051618099212646 sec
7.922858238220215 sec

看起来问题不在于DataLoader，而在于dataset。当我在第一次“长”迭代之后删除并重新初始化DataLoader对象时，一切仍然正常。当我重新初始化数据集时——长的第一次迭代再次出现。 此外，在这个时期，我通过htop跟踪我的cpu利用率，将num_workers设置为32，在第一个时期，利用率非常低；32个核中只有1-2个在工作，在其他时代~所有核都在工作

Answer 1

斯拉夫卡

我没有下载整个GLR2020数据集，但我能够在我本地拥有的图像数据集上观察到这种效果（80000张jpg图像，大小约为400x400）

为了找出性能差异的原因，我尝试了以下方法：

1. 将扩充减少为仅调整大小
2. 只测试ImgDataset.__getitem__()函数
3. ImgDataset.__getitem__()无增广
4. 只需加载原始jpg图像并将其从数据集中传递出去，而无需进行numpy转换

结果表明，差异来自图像加载时间。Python（或OS本身）实现了某种缓存，在下面的测试中多次加载映像时会观察到这种缓存

for i in range(5):    
    t0 = time.time()
    data = cv2.imread(filename)
    print (time.time() - t0)
    
0.03395271301269531
0.0010004043579101562
0.0010004043579101562
0.0010008811950683594
0.001001119613647461

当只从文件读取到变量时，也会观察到同样的情况

for i in range(5):    
    t0 = time.time()
    with open(filename, mode='rb') as file: 
        data = file.read()
    print (time.time() - t0)

0.036234378814697266
0.0028831958770751953
0.0020024776458740234
0.0031833648681640625
0.0028734207153320312

降低加载速度的一种方法是将数据保持在非常快的本地SSD上。若大小允许，尝试将数据集的一部分加载到RAM中，并从那个里编写自定义数据加载器来馈送

顺便说一句，根据我的发现，这种效果应该可以在任何数据集上重现——看看是否使用了不同的驱动器或缓存

Answer 2

操作系统似乎正在缓存对数据集的IO访问。要检查这是否确实是问题所在，请尝试在第一个历元之后运行sync; echo 3 > /proc/sys/vm/drop_caches（在Ubuntu上）。如果执行此操作时第二个历元同样慢，则是缓存使后续读取速度快得多

如果您使用的是HDD，那么通过将所有的小图像文件放在磁盘上，您可能会在第一个历元中获得显著的速度提升

您可以使用SquashFS（它是随Ubuntu预装的）将整个数据集压缩到单个文件中，然后将该文件作为目录装入，并像以前一样访问它（除了现在图像位于磁盘上）。装载的目录是只读的

例如

mksquashfs /path/to/data data.sqsh
mount data.sqsh /path/to/data_sqsh -t squashfs -o loop

然后您可以使用/path/to/data_sqsh，使用方式与使用/path/to/data完全相同。重新启动计算机时，您必须重新安装它

见：https://tldp.org/HOWTO/SquashFS-HOWTO/creatingandusing.html