Pytorch DataLoader primera época extremadamente lenta

Question

Cuando creo un PyTorch DataLoader y empiezo a iterar, obtengo una primera época extremadamente lenta (x10 - x30 más lenta que todas las siguientes épocas). Además, este problema ocurre solo con el conjunto de datos de trenes del reconocimiento de hitos de Google 2020 de Kaggle. No puedo reproducir esto en imágenes sintéticas, también intenté crear una carpeta con 500k imágenes de GLR2020 y todo funcionó bien. Encontré algunos problemas similares en el foro de PyTorch sin ninguna solución.

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

Aquí hay algunos ejemplos de velocidad de ejecución con diferentes num_workersen 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 el problema no es con DataLoader, sino con el conjunto de datos. Cuando elimino y reinicializo el objeto DataLoader después de la primera iteración "larga", todo sigue funcionando bien. Cuando reinicializo el conjunto de datos, vuelve a aparecer la primera iteración larga. Además, realicé un seguimiento de la utilización de mi CPU htopdurante estas épocas con el valor num_workers32, y durante la primera época, la utilización es muy baja; sólo 1-2 de 32 núcleos están funcionando, durante otras épocas ~ todos los núcleos están funcionando.

Answer 1

Slavka,

No descargué todo el conjunto de datos GLR2020, pero pude observar este efecto en el conjunto de datos de imágenes que tenía localmente (80000 imágenes jpg de aproximadamente 400x400 tamaño).

Para encontrar las razones de la diferencia en el rendimiento, intenté lo siguiente:

1. reducir el aumento a solo cambiar el tamaño
2. probando solo la ImgDataset.__getitem__()función
3. ImgDataset.__getitem__() sin aumento
4. simplemente cargando la imagen jpg sin procesar y pasándola desde el conjunto de datos sin siquiera una conversión numpy.

Resulta que la diferencia proviene del tiempo de carga de la imagen. Python (o el propio sistema operativo) implementa algún tipo de almacenamiento en caché que se observa al cargar la imagen varias veces en la siguiente prueba.

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

Lo mismo se observa cuando solo se lee de archivo a variable

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

Una forma de reducir la velocidad de carga es mantener los datos en un SSD local muy rápido. Si el tamaño lo permite, intente cargar parte del conjunto de datos en la RAM y escribir un cargador de datos personalizado para alimentar desde allí ...

Por cierto, según mis hallazgos, este efecto debería ser reproducible con cualquier conjunto de datos; vea si usó diferentes unidades o algún almacenamiento en caché.

Answer 2

Parece que el sistema operativo está almacenando en caché el acceso de E / S al conjunto de datos. Para comprobar si este es definitivamente el problema, intente ejecutar sync; echo 3 > /proc/sys/vm/drop_caches(en Ubuntu) después de la primera época. Si la segunda época es igualmente lenta cuando hace esto, entonces es el almacenamiento en caché lo que hace que las lecturas posteriores sean mucho más rápidas.

Si está utilizando un disco duro, puede obtener mejoras de velocidad significativas para su primera época al colocar todos sus archivos de imagen pequeños en el disco.

Puede usar SquashFS (viene preinstalado con Ubuntu) para comprimir todo su conjunto de datos en un solo archivo, luego montar ese archivo como un directorio y acceder a él como antes (excepto que ahora las imágenes están ubicadas en el disco). El directorio montado es de solo lectura.

p.ej

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

Entonces puede usar /path/to/data_sqshexactamente de la misma manera que usó /path/to/data. Tendrá que volver a montarlo cuando reinicie su computadora

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