Pytorch 默认数据加载器因大型图像分类训练集而卡住

Question

我正在 Pytorch 中训练图像分类模型，并使用其默认数据加载器来加载我的训练数据。我有一个非常大的训练数据集，因此每个类通常有几千个样本图像。过去我已经用大约 20 万张图像训练过模型，没有出现任何问题。然而我发现，当图像总数超过一百万张时，Pytorch 数据加载器就会卡住。

\n\n

我相信当我打电话时代码挂起datasets.ImageFolder(...)。当我按下 Ctrl-C 时，输出始终如下：

\n\n

Traceback (most recent call last):                                                                                                 \xe2\x94\x82\n  File "main.py", line 412, in <module>                                                                                            \xe2\x94\x82\n    main()                                                                                                                         \xe2\x94\x82\n  File "main.py", line 122, in main                                                                                                \xe2\x94\x82\n    run_training(args.group, args.num_classes)                                                                                     \xe2\x94\x82\n  File "main.py", line 203, in run_training                                                                                        \xe2\x94\x82\n    train_loader = create_dataloader(traindir, tfm.train_trans, shuffle=True)                                                      \xe2\x94\x82\n  File "main.py", line 236, in create_dataloader                                                                                   \xe2\x94\x82\n    dataset = datasets.ImageFolder(directory, trans)                                                                               \xe2\x94\x82\n  File "/home/username/.local/lib/python3.5/site-packages/torchvision/datasets/folder.py", line 209, in __init__     \xe2\x94\x82\n    is_valid_file=is_valid_file)                                                                                                   \xe2\x94\x82\n  File "/home/username/.local/lib/python3.5/site-packages/torchvision/datasets/folder.py", line 94, in __init__      \xe2\x94\x82\n    samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file)                                                     \xe2\x94\x82\n  File "/home/username/.local/lib/python3.5/site-packages/torchvision/datasets/folder.py", line 47, in make_dataset  \xe2\x94\x82\n    for root, _, fnames in sorted(os.walk(d)):                                                                                     \xe2\x94\x82\n  File "/usr/lib/python3.5/os.py", line 380, in walk                                                                               \xe2\x94\x82\n    is_dir = entry.is_dir()                                                                                                        \xe2\x94\x82\nKeyboard Interrupt                                                                                                                       \n

\n\n

我认为某处可能存在死锁，但是根据 Ctrl-C 的堆栈输出，它看起来不像在等待锁。所以后来我认为数据加载器很慢，因为我试图加载更多的数据。我让它运行了大约 2 天，但没有取得任何进展，在加载的最后 2 小时中，我检查了 RAM 使用量保持不变。过去，我还能够在不到几个小时的时间内加载包含超过 20 万张图像的训练数据集。我还尝试将我的 GCP 机器升级为 32 个核心、4 个 GPU 和超过 100GB 的 RAM，但似乎在加载一定量的内存后，数据加载器就会卡住。

\n\n

我很困惑数据加载器在循环目录时如何卡住，并且我仍然不确定它是否卡住或只是非常慢。有什么方法可以更改 Pytortch 数据加载器，使其能够处理超过 100 万张图像进行训练？任何调试建议也将受到赞赏！

\n\n

谢谢你！

\n

Answer 1

这不是 的问题DataLoader，而是它的工作原理的问题torchvision.datasets.ImageFolder（以及为什么数据越多它的工作效果就越差）。

它挂在这条线上，如您的错误所示：

for root, _, fnames in sorted(os.walk(d)): 

来源可以在这里找到。

根本问题是它在 Giant 中保留了每个path和 对应项，请参阅下面的代码（为了简洁起见，删除了一些内容）：labellist

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

显然，图像将包含 100 万个字符串（也相当长），并且对应int于类别，这绝对是很多并且取决于 RAM 和 CPU。

不过，您可以创建自己的数据集（前提是您事先更改图像的名称），因此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的指定类、循环执行或其他方式。

最后，torch.utils.data.DataLoader按平常使用，例如：

dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)