pytorch DataLoader 极慢的第一个纪元

Question

当我创建一个 PyTorch DataLoader 并开始迭代时——我得到一个非常慢的第一个纪元（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，而在于数据集。当我在第一次“长”迭代后删除并重新初始化 DataLoader 对象时，一切仍然正常。当我重新初始化数据集时——长第一次迭代再次出现。此外，在 num_workers 设置为 32 的这个时期，我通过 htop 跟踪了我的 cpu 利用率，并且在第一个时期，利用率非常低； 32 个核心中只有 1-2 个在工作，在其他时期〜所有核心都在工作。

Answer 1

斯拉夫卡

TLDR：这是缓存效果。

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

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

将增强减少到仅调整大小
仅测试 ImgDataset.__getitem__() 函数
ImgDataset.__getitem__() 无增强
只需加载原始 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

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

如果您使用的是 HDD，那么您可以通过 co-locating 磁盘上的所有小图像文件在您的第一个时期获得显着的速度提升。

您可以使用 SquashFS（pre-installed 和 Ubuntu 附带）将整个数据集压缩到单个文件中，然后将该文件挂载为目录并像以前一样访问它（除了现在图像在磁盘上 co-located）。挂载目录为read-only.

例如

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

然后您就可以像使用 /path/to/data 一样使用 /path/to/data_sqsh。当您重新启动计算机时，您将必须 re-mount 它

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

pytorch DataLoader 极慢的第一个纪元

pytorch DataLoader extremely slow first epoch

python

multiprocessing

pytorch

dataloader