pytorch DataLoader 第一个时代非常慢

Question

当我创建一个 PyTorch DataLoader 并开始迭代时——我得到了一个非常慢的第一个纪元（x10--x30 比所有下一个纪元都慢）。此外，此问题仅出现在来自 Kaggle 的 Google 地标识别 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')

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

#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 对象时，一切仍然正常。当我重新初始化数据集时——长的第一次迭代再次出现。此外，我通过设置为 32 在此 epochhtop期间跟踪我的 CPU 利用率num_workers，并且在第一个 epoch 期间，利用率非常低；32 个核心中只有 1-2 个在工作，在其他时期 ~ 所有核心都在工作。

Answer 1

斯拉夫卡，

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

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

1. 将增强减少到仅调整大小
2. 测试只是ImgDataset.__getitem__()功能
3. ImgDataset.__getitem__() 没有增强
4. 只需加载原始 jpg 图像并从数据集中传递它，甚至无需进行 numpy 转换。

事实证明，差异来自图像加载时间。Python（或操作系统本身）实现了某种缓存，在以下测试中多次加载图像时会观察到这种缓存。

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 中并编写自定义数据加载器以从那里输入...

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