Load the entire dataset into memory

[ScanFun]

🚀 The feature

I plan to add a feature for loading datasets into memory to the built-in datasets in the datasets module in torchvision. Maybe we can add a parameter to the initialization function of the datasets, such as "to_memory"?

I have now implemented the prototype of this feature by modifying ImageFolder:

# ATTENTION: This code is UNFINISHED!
def __init__(
    self,
    root: str,
    loader: Callable[[str], Any],
    extensions: Optional[Tuple[str, ...]] = None,
    transform: Optional[Callable] = None,
    target_transform: Optional[Callable] = None,
    is_valid_file: Optional[Callable[[str], bool]] = None,
) -> None:
    super().__init__(root, transform=transform, target_transform=target_transform)
    classes, class_to_idx = self.find_classes(self.root)
    samples = self.make_dataset(self.root, class_to_idx, extensions, is_valid_file)

    self.loader = loader
    self.extensions = extensions

    self.classes = classes
    self.class_to_idx = class_to_idx
    self.samples = samples
    self.targets = [s[1] for s in samples]
    self.images = []

    for s in samples:
        sample = self.loader(s[0])
        if self.transform is not None:
            sample = self.transform(sample)
        self.images.append(sample)

def __getitem__(self, index: int) -> Tuple[Any, Any]:
    """
    Args:
        index (int): Index

    Returns:
        tuple: (sample, target) where target is class_index of the target class.
    """
    path, target = self.samples[index]
    sample = self.images[index]
    if self.target_transform is not None:
        target = self.target_transform(target)
    return sample, target

Motivation, pitch

I found that training was unusually slow when I was using a larger dataset for one of my projects (I was using ImageFolder to build my dataset).
After looking at the data given by pytorch_profiler, I found that aten::copy_ takes up most of the time.
I thought this meant I had an IO bottleneck, so I wanted to see if I could load the entire dataset into memory, but torchvision didn't seem to support it.

[abhi-glitchhg]

According to me loading the entire dataset into memory is not an ideal way.

I thought this meant I had an IO bottleneck, so I wanted to see if I could load the entire dataset into memory, but torchvision didn't seem to support it.

In such cases, i increase the number of workers with the dataloader. It speeds up the data loading part. Have you tried that?

[NicolasHug]

Hi @ScanFun , thanks for the feature request.

    for s in samples:
        sample = self.loader(s[0])
        if self.transform is not None:
            sample = self.transform(sample)
        self.images.append(sample)

you probably don't really want to do that as this is caching the images after they've been transformed: you would get the same transformation for a given image across epochs, and that could lead to less efficient augmentation / training / accuracy.

---

The difficulty with providing a built-in caching mechanism is that there may be more than one place where users would want caching to happen: typically either a) before decoding or b) after decoding, before transforms, c) during transforms (but before any random transform).

Perhaps the best way to address this in torchvision is to essentially... not address it all and let users implement their own custom wrapper classes? It should be fairly easy to implement any of the above option by wrapping around ImageFolder(). (For other built-in dataset classes though, it may be a bit more tricky and involve monkey patching).

Any thought @pmeier ?

[ScanFun]

According to me loading the entire dataset into memory is not an ideal way.

I thought this meant I had an IO bottleneck, so I wanted to see if I could load the entire dataset into memory, but torchvision didn't seem to support it.

In such cases, i increase the number of workers with the dataloader. It speeds up the data loading part. Have you tried that?

Thank you for your reply!
Well, I have already added the num of the Dataloade's num_workers have reached the maximum allowed on my computer.

So I can't find a better solution to it than upgrading my hardware.

[ScanFun]

@NicolasHug
Wow, thank you for your reply. But I don't seem to understand the following sentence:

The difficulty with providing a built-in caching mechanism is that there may be more than one place where users would want caching to happen: typically either a) before decoding or b) after decoding, before transforms, c) during transforms (but before any random transform).

What you mean is that if we do a built-in caching mechanism, we must establish a caching mechanism for the original data, decoded data, and transformed data?

Also, does Pytorch re-init the dataset I use every epoch? Is the problem with my code that I cache data that has been transformed?
(Because the original transform processing statement was in the getitem(), Pytorch repeatedly called getitem whenever needed, and also repeated the transform?）

Do I need to move the statement for the transform, like this:

#Remove sample = self.transform(sample)) from init()
def __getitem__(self, index: int) -> Tuple[Any, Any]:
    """
    Args:
        index (int): Index

    Returns:
        tuple: (sample, target) where target is class_index of the target class.
    """
    path, target = self.samples[index]
    sample = self.images[index]
    if self.transform is not None:
        sample = self.transform(sample)
    if self.target_transform is not None:
        target = self.target_transform(target)
    return sample, target

[NicolasHug]

What you mean is that if we do a built-in caching mechanism, we must establish a caching mechanism for the original data, decoded data, and transformed data?

You as a user don't need to cache in all these places. What I meant is that some users may want to cache in a), others in b), others in c). And supporting all options in torchvision, along with cache-eviction / cache size mechanisms is something difficult /complex to do. But implementing one single specific use-case is simple-enough to do for users.

Also, does Pytorch re-init the dataset I use every epoch?

The torch.utils.data.DataLoader won't re-instantiate the dataset (you already pass an instance), it will only access it through the __getitem__ method. So if you call dataset[idx] twice (whether across epochs or not), depending on what you cache, you may end up with the same transformed img (which is not desirable).

Is the problem with my code that I cache data that has been transformed? (Because the original transform processing statement was in the getitem(), Pytorch repeatedly called getitem whenever needed, and also repeated the transform?）

Yes (assuming your cache is self.images)

Do I need to move the statement for the transform, like this:

Yes, that's implementing b) i.e. you're caching decoded images. Your cache can potentially grow quite big but if your dataset is small enough that may work. If you only care about optimizing for io latency, option a) might be better, i.e. cache the bytes before they're decoded (e.g. cache BytesIO objects or pure bytes tensors, stuff like that)

[ScanFun]

@NicolasHug
Thank you for your reply!

---

You as a user don't need to cache in all these places. What I meant is that some users may want to cache in a), others in b), others in c). And supporting all options in torchvision, along with cache-eviction / cache size mechanisms is something difficult /complex to do. But implementing one single specific use-case is simple-enough to do for users.

Sorry, I actually don't know under what circumstances the user wants to cache other than case b.

For example:
1.If we cache the data set when case b, then the data is still cached when case c (because the cached data when case b is already in memory).
2. As for caching when case A, I don't know why users need to put undecoded raw data into memory during training.
3. And if it is cached when c, as you said, it will reduce the accuracy and so on.

Maybe "to_memory" is causing confusion, maybe I should change its name. . .
My intention is not to put the entire dataset and all the data derived from it into memory, I do this to speed up training.

---

The torch.utils.data.DataLoader won't re-instantiate the dataset (you already pass an instance), it will only access it through the getitem method. So if you call dataset[idx] twice (whether across epochs or not), depending on what you cache, you may end up with the same transformed img (which is not desirable).

Thank you very much for pointing this out!
I now realize that there will be problems if I put the transform statement in init().

---

Yes, that's implementing b) i.e. you're caching decoded images. Your cache can potentially grow quite big but if your dataset is small enough that may work. If you only care about optimizing for io latency, option a) might be better, i.e. cache the bytes before they're decoded (e.g. cache BytesIO objects or pure bytes tensors, stuff like that)

If we decode in getitem(), wouldn't it increase the time each time we get the data?
(well, I don't actually know the impact of loader() on time)

[NicolasHug]

1.If we cache the data set when case b, then the data is still cached when case c (because the cached data when case b is already in memory).

Yes although see reply to point 3) below

2. As for caching when case A, I don't know why users need to put undecoded raw data into memory during training.

To avoid io/filesystem latency. If you have a fast fs it's OK, but if you have an old hard-drive, simply physically accessing and reading a single jpeg can be the bottleneck. That's why people use super fast file systems, or store huge blobs of data containing multiple images instead of individual jpeg files. (all datasets in torchvision access individual image files though).

3. And if it is cached when c, as you said, it will reduce the accuracy and so on.

Not if you only cache before any random transforms, e.g. if your first transform is to resize to a fixed size, you could cache resized images instead of caching non-resized images. That would gain you some speed (and depending on the original size of the images, it could also save you some memory!). Just don't cache images that have been transformed with RNG involved.

If we decode in getitem(), wouldn't it increase the time each time we get the data?

Yes. But that would reduce the cache size because you'd be storing raw compressed jpegs instead of decoded ones. It's almost always about the speed vs memory tradeoff.

[ScanFun]

@NicolasHug Sorry for taking so long time to reply.

To avoid io/filesystem latency. If you have a fast fs it's OK, but if you have an old hard-drive, simply physically accessing and reading a single jpeg can be the bottleneck. That's why people use super fast file systems, or store huge blobs of data containing multiple images instead of individual jpeg files. (all datasets in torchvision access individual image files though).

Yes, but if this feature is only intended to improve the training speed of certain users, we don't seem to need to consider caching the raw dataset.

Not if you only cache before any random transforms, e.g. if your first transform is to resize to a fixed size, you could cache resized images instead of caching non-resized images. That would gain you some speed (and depending on the original size of the images, it could also save you some memory!). Just don't cache images that have been transformed with RNG involved.

So, um, this feature was not actually optimized for specific situations.

Essentially, the improvement that can be achieved with this feature is very limited (such as if you don't need this feature if you have a fast FS, etc.).

Moreover, if we optimize for all possible usage scenarios, this feature will be very complex.

In fact, only caching the decoded data can meet the needs of most training-model-only users.

Regarding the issue of occupying too much memory, for example, if the user's memory is not large enough and FS is not fast enough, they should upgrade their hardware or seek other solutions.

Yes. But that would reduce the cache size because you'd be storing raw compressed jpegs instead of decoded ones. It's almost always about the speed vs memory tradeoff.

Just like the second reply, make this feature simple.

If we specialize in certain optimizations, it may significantly increase code complexity.

[ScanFun]

Um. Well, can anyone tell me how many days it usually takes to discuss the issue of new features?

[NicolasHug]

Above I shared insights about the different places where users may want caching to happen. This is just the tip of the iceberg because of course, on top of where caching happens, one needs to worry about other details like cache-eviction mechanisms, cache size, etc. The point is: supporting caching as we would need to support it in torchvision isn't something easy, and we don't really have plans to do that right now.

OTOH, implementing the specific caching strategy that you need is rather easy and can be handled on a adhoc basis by users, so we recommend that you just implement that in your own code.