pytorch DataLoader extremamente lento na primeira época

Question

Quando eu crio um PyTorch DataLoader e começo a iterar - obtenho uma primeira época extremamente lenta (x10 - x30 mais lenta do que todas as próximas épocas). Além disso, esse problema ocorre apenas com o conjunto de dados de trem do reconhecimento de referência do Google 2020 de Kaggle. Não consigo reproduzir em imagens sintéticas, também, tentei criar uma pasta com 500k de imagens do GLR2020, e deu tudo certo. Encontrou poucos problemas semelhantes no fórum PyTorch sem nenhuma solução.

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')

Aqui estão alguns exemplos de velocidade de execução com diferentes num_workersem 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

Parece que o problema não é com DataLoader, mas com dataset. Quando eu excluo e reinicializo o objeto DataLoader após a primeira iteração "longa", tudo ainda funciona bem. Quando eu reinicializo o conjunto de dados - a primeira iteração longa aparece novamente. Além disso, acompanhei a utilização da minha cpu por meio htopdessas épocas com num_workersdefinido para 32 e, durante a primeira época, a utilização é muito baixa; apenas 1-2 de 32 núcleos estão funcionando, durante outras épocas ~ todos os núcleos estão funcionando.

Answer 1

Slavka,

Eu não baixei o conjunto de dados GLR2020 completo, mas pude observar esse efeito no conjunto de dados de imagens que eu tinha localmente (80000 imagens jpg de tamanho aproximado de 400x400).

Para encontrar os motivos da diferença de desempenho, tentei o seguinte:

1. reduzindo o aumento para apenas redimensionar
2. testando apenas a ImgDataset.__getitem__()função
3. ImgDataset.__getitem__() sem aumento
4. basta carregar a imagem jpg bruta e passá-la do conjunto de dados sem nem mesmo fazer uma conversão numpy.

Acontece que a diferença vem do tempo de carregamento da imagem. Python (ou o próprio sistema operacional) implementa algum tipo de cache que é observado ao carregar a imagem várias vezes no teste a seguir.

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

o mesmo é observado quando apenas lendo de arquivo para variável

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

Uma maneira de reduzir a velocidade de carregamento é manter os dados em um SSD local muito rápido. Se o tamanho permitir, tente carregar parte do conjunto de dados na RAM e escrever um dataloader personalizado para alimentar a partir daí ...

BTW Com base em minhas descobertas, este efeito deve ser reproduzível com qualquer conjunto de dados - veja se você usou drives diferentes ou algum cache.

Answer 2

Parece que o sistema operacional está armazenando em cache o acesso IO ao conjunto de dados. Para verificar se este é definitivamente o problema, tente executar sync; echo 3 > /proc/sys/vm/drop_caches(no Ubuntu) após a primeira época. Se a segunda época for igualmente lenta quando você fizer isso, é o armazenamento em cache que está tornando as leituras subsequentes muito mais rápidas.

Se você estiver usando um HDD, poderá obter melhorias de velocidade significativas para a sua primeira época, colocando todos os seus pequenos arquivos de imagem no disco.

Você pode usar o SquashFS (ele vem pré-instalado com o Ubuntu) para compactar todo o conjunto de dados em um único arquivo, depois montar esse arquivo como um diretório e acessá-lo como antes (exceto que agora as imagens estão co-localizadas no disco). O diretório montado é somente leitura.

por exemplo

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

Então você pode usar /path/to/data_sqshexatamente da mesma maneira que você usou /path/to/data. Você terá que remontá-lo ao reiniciar o computador

Vejo: https://tldp.org/HOWTO/SquashFS-HOWTO/creatingandusing.html