Это не проблема с DataLoader, это проблема с torchvision.datasets.ImageFolder и с тем, как он работает (и почему он работает намного хуже, чем больше у вас данных).
Он висит на этой строке, как указывает ваша ошибка:
for root, _, fnames in sorted(os.walk(d)):
Источник можно найти здесь .
Основная проблема заключается в том, что он сохраняет каждый path и соответствующий label в гиганте list, см. Код ниже (несколько вещей для краткости исключены):
def make_dataset(dir, class_to_idx, extensions=None, is_valid_file=None):
images = []
dir = os.path.expanduser(dir)
# Iterate over all subfolders which were found previously
for target in sorted(class_to_idx.keys()):
d = os.path.join(dir, target) # Create path to this subfolder
# Assuming it is directory (which usually is the case)
for root, _, fnames in sorted(os.walk(d, followlinks=True)):
# Iterate over ALL files in this subdirectory
for fname in sorted(fnames):
path = os.path.join(root, fname)
# Assuming it is correctly recognized as image file
item = (path, class_to_idx[target])
# Add to path with all images
images.append(item)
return images
Очевидно, что изображения будут содержать 1 миллион строк (также довольно длинных) и соответствующие int для классов, которые определенно являются много и зависит от ОЗУ и ЦП.
Вы можете создавать свои собственные наборы данных (при условии, что вы заранее изменяете имена своих изображений), поэтому память не будет занята dataset.
Настройка структуры данных
Структура вашей папки должна выглядеть следующим образом:
root
class1
class2
class3
...
Использовать, сколько классов вам нужно / нужно.
Теперь каждый class должны иметь следующие данные:
class1
0.png
1.png
2.png
...
Учитывая, что вы можно перейти к созданию наборов данных.
Создать наборы данных
Ниже torch.utils.data.Dataset использует PIL для открытия изображений, вы можете сделать это по-другому, хотя:
import os
import pathlib
import torch
from PIL import Image
class ImageDataset(torch.utils.data.Dataset):
def __init__(self, root: str, folder: str, klass: int, extension: str = "png"):
self._data = pathlib.Path(root) / folder
self.klass = klass
self.extension = extension
# Only calculate once how many files are in this folder
# Could be passed as argument if you precalculate it somehow
# e.g. ls | wc -l on Linux
self._length = sum(1 for entry in os.listdir(self._data))
def __len__(self):
# No need to recalculate this value every time
return self._length
def __getitem__(self, index):
# images always follow [0, n-1], so you access them directly
return Image.open(self._data / "{}.{}".format(str(index), self.extension))
Теперь вы можете легко создавать свои наборы данных (структура папок, как указано выше:
root = "/path/to/root/with/images"
dataset = (
ImageDataset(root, "class0", 0)
+ ImageDataset(root, "class1", 1)
+ ImageDataset(root, "class2", 2)
)
Вы можете добавить столько datasets с указанными классами, сколько вы будете sh, сделайте это в l oop или что-то еще.
Наконец, используйте torch.utils.data.DataLoader как обычно, например:
dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)