如何加速ImageNet的"ImageFolder“

Question

我在一所大学，所有的文件系统都在远程系统中，无论我用我的账号登录到哪里，我总是可以访问我的主目录。即使我通过SSH命令登录到GPU服务器。这就是我使用GPU服务器读取数据的条件。

目前，我使用PyTorch在ImageNet上从头开始训练ResNet，我的代码只使用同一台计算机上的所有GPU，我发现"torchvision.datasets.ImageFolder“几乎需要两个小时。

你能提供一些关于如何加速torchvision.datasets.ImageFolder的经验吗？非常感谢。

Answer 1

为什么要花这么长时间？

设置ImageFolder可能需要很长时间，尤其是当图像存储在速度较慢的远程磁盘上时。出现这种延迟的原因是数据集的__init__函数会遍历图像文件夹中的所有文件，并检查该文件是否为图像文件。对于ImageNet来说，这可能需要相当长的时间，因为有100多万个文件需要检查。

你能做什么？

• 正如Kevin Sun已经指出的，将数据集复制到本地(可能更快)存储可以显著提高速度您可以创建一个修改后的数据集类，它不会读取所有文件，但依赖于文件的缓存列表-一个只需提前准备一次并用于所有运行的缓存列表。

Answer 2

如果您确定文件夹结构不会更改，则可以使用以下命令缓存该结构(而不是太大的数据)：

import json
from functools import wraps
from torchvision.datasets import ImageNet

def file_cache(filename):
    """Decorator to cache the output of a function to disk."""
    def decorator(f):
        @wraps(f)
        def decorated(self, directory, *args, **kwargs):
            filepath = Path(directory) / filename
            if filepath.is_file():
                out = json.loads(filepath.read_text())
            else:
                out = f(self, directory, *args, **kwargs)
                filepath.write_text(json.dumps(out))
            return out
        return decorated
    return decorator

class CachedImageNet(ImageNet):
    @file_cache(filename="cached_classes.json")
    def find_classes(self, directory, *args, **kwargs):
        classes = super().find_classes(directory, *args, **kwargs)
        return classes

    @file_cache(filename="cached_structure.json")
    def make_dataset(self, directory, *args, **kwargs):
        dataset = super().make_dataset(directory, *args, **kwargs)
        return dataset