# Best practice to cache the entire dataset during first epoch

**URL:** <https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608>\
**Category:** Uncategorized\
**Created:** [June 13, 2018, 6:52am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608 "2018-06-13T06:52:05Z")\
**Posts on this page:** 13\
**Page:** 1

<div class="post-metadata">

**Author:** ![bowenc0221](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bowenc0221/32/4912_2.png) [@bowenc0221](https://discuss.pytorch.org/u/bowenc0221)\
**Post date:** [June 13, 2018, 6:52am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/1 "2018-06-13T06:52:06Z")

</div>

Hi,

Currently, I am in a situation: the dataset is stored in a single file on a shared file system and too many processes accessing the file will cause a slow down to the file system (for example, 40 jobs each with 20 workers will end up 800 processes reading from the same file). So I plan to load the dataset to the memory.

I have enough memory (~500G) to hold the entire dataset (for example, ImageNet 1k), but loading the dataset before training is too slow. I would like to know if there is a good way to cache the entire dataset during the first epoch so that after first epoch workers will close the file and read directly from memory.

Thanks.

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [June 13, 2018, 11:06am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/2 "2018-06-13T11:06:46Z")

</div>

One approach would be to store the already loaded data in the `Dataset` and use it afterwards.  
The downside is that you are limited to `num_workers=0` in the first epoch.

```auto
class MyDataset(Dataset):
    def __init__ (self, use_cache=False):
        self.data = torch.randn(100, 1)
        self.cached_data = []
        self.use_cache = use_cache
        
    def __getitem__ (self, index):
        if not self.use_cache:
            x = self.data[index] # your slow data loading
            self.cached_data.append(x)
        else:
            x = self.cached_data[index]
        
        return x
    
    def set_use_cache(self, use_cache):
        if use_cache:
            self.cached_data = torch.stack(self.cached_data)
        else:
            self.cached_data = []
        self.use_cache = use_cache
    
    def __len__ (self):
        return len(self.data)

dataset = MyDataset(use_cache=False)
loader = DataLoader(
    dataset,
    num_workers=0,
    shuffle=False
)

for data in loader:
    print(len(loader.dataset.cached_data))

loader.dataset.set_use_cache(use_cache=True)
loader.num_workers = 2
for data in loader:
    print(len(loader.dataset.cached_data))

```

---

<div class="post-metadata">

**Author:** ![Sagar\_Bhonde](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sagar_bhonde/32/17288_2.png) [@Sagar\_Bhonde](https://discuss.pytorch.org/u/Sagar_Bhonde)\
**Post date:** [November 4, 2019, 6:10pm UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/3 "2019-11-04T18:10:07Z")

</div>

In order to use this for images-

I tried this solution for images to which I got this error

File “/home/sagarbhonde/kiwi/kd/dataset.py”, line 90, in set\_use\_cache  
self.cached\_images = torch.stack(self.cached\_images)  
TypeError: stack(): argument ‘tensors’ (position 1) must be tuple of Tensors, not Tensor

So you might want to change `set_use_cache` method to:

```auto
    def set_use_cache(self, use_cache):
        if use_cache:
            x_img = tuple(self.cached_images)
            self.cached_images = torch.stack(x_img)
        else:
            self.cached_images = []
        self.use_cache = use_cache

```

---

<div class="post-metadata">

**Author:** ![bfeeny](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bfeeny/32/26059_2.png) [@bfeeny](https://discuss.pytorch.org/u/bfeeny)\
**Post date:** [July 29, 2020, 5:58pm UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/4 "2020-07-29T17:58:05Z")

</div>

@ptrblck I was just reading this response you wrote back in 2018. I am trying to do some basic speed up on my Dataloader (loading images). Is what you wrote here still valid? I saw you also responded to a similar post [here](https://discuss.pytorch.org/t/dataloader-resets-dataset-state/27960/3) and I wasn’t sure if you have to use multiprocessing Array in most cases. I liked the simplicity of what you show here. I am doing 5 folds, 15 epochs per fold, and Test Time Augmentation x 15 on Validation and Test. So that’s a lot of time reading the same files from disk over and over, and I am hoping I can use something like this to speed things up a bit.

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [July 30, 2020, 3:44am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/5 "2020-07-30T03:44:00Z")

</div>

The posted approach here might still be valid, but note that the first pass is using `num_workers=0` in order to add the loaded samples into the `Dataset`.  
The following epochs might use more workers, where each process will create a copy of the complete `Dataset`, which would increase the memory usage by `num_workers`.

The linked post uses the shared arrays to directly paste the loaded samples into it using multiple workers.

Let me know, if you encounter any issues.

---

<div class="post-metadata">

**Author:** ![bfeeny](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bfeeny/32/26059_2.png) [@bfeeny](https://discuss.pytorch.org/u/bfeeny)\
**Post date:** [July 30, 2020, 3:47am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/6 "2020-07-30T03:47:46Z")

</div>

@ptrblck I did encounter some issues, I’ll post in the other thread though as its relevant to that method.

---

<div class="post-metadata">

**Author:** ![Roye](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/4da419/32.png) [@Roye](https://discuss.pytorch.org/u/Roye)\
**Post date:** [April 16, 2021, 2:31am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/7 "2021-04-16T02:31:51Z")

</div>

Thank you for your answer, it works great.  
Can you explain why num\_workers must be zero, and what happens when it is not?

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [April 16, 2021, 5:05am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/8 "2021-04-16T05:05:30Z")

</div>

The first epoch would fill the “cache” in the original `Dataset` object using a single worker. The other epochs would then use multiple workers and reuse this cache, since each worker would create a copy of the dataset.  
Note that this is not an optimal approach, but more a proof of concept.

---

<div class="post-metadata">

**Author:** ![Roye](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/4da419/32.png) [@Roye](https://discuss.pytorch.org/u/Roye)\
**Post date:** [April 16, 2021, 1:23pm UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/9 "2021-04-16T13:23:20Z")

</div>

I placed the cache object within the dataset class as a class variable, so copies would still “see and use” the same cached object. Although this works on the console, when using num\_workers \> 0 it doesn’t, and copying shouldn’t cause the effect ( as much as I understanding python’s internals, and I guess I don’t ).

in your reply, you wrote: "not an optimal approach… ", what do you recommend when handling images as files?

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [April 17, 2021, 5:42am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/10 "2021-04-17T05:42:50Z")

</div>

You could use shared arrays instead of this list, which forces you to use a the main process while the cache is being filled.

---

<div class="post-metadata">

**Author:** ![Xpitfire](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/xpitfire/32/38191_2.png) [@Xpitfire](https://discuss.pytorch.org/u/Xpitfire)\
**Post date:** [May 12, 2021, 3:00pm UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/11 "2021-05-12T15:00:43Z")

</div>

I tried to compile a version I used to cache my preprocessed images. This impl. saves some memory and gives some speedup since it avoids resizing the images in each process. About the `torch.multiprocessing.set_sharing_strategy('file_system')` call I am not absolutely sure and maybe @ptrblck has some comments on this, but it was suggested by the PyTorch framework and without it I got an exception after 14+ hours of training when passing on the same dataloader to 30+ different model training loops.

```python
import torch
# used at the beginning of your program
torch.multiprocessing.set_sharing_strategy('file_system')
import torchvision.transforms as transforms
from PIL import Image
from multiprocessing import Manager
from torch.utils.data import Dataset, Dataloader

class DatasetCache(object):
    def __init__ (self, manager, use_cache=True):
        self.use_cache = use_cache
        self.manager = manager
        self._dict = manager.dict()

    def is_cached(self, key):
        if not self.use_cache:
            return False
        return str(key) in self._dict

    def reset(self):
        self._dict.clear()

    def get(self, key):
        if not self.use_cache:
            raise AttributeError('Data caching is disabled and get funciton is unavailable! Check your config.')
        return self._dict[str(key)]

    def cache(self, key, img, lbl):
        # only store if full data in memory is enabled
        if not self.use_cache:
            return
        # only store if not already cached
        if str(key) in self._dict:
            return
        self._dict[str(key)] = (img, lbl)

class CachedListDataset(Dataset):
    '''Load image/labels from a list file.
    The list file is like:
      a.jpg label ...
    '''
    def __init__ (self, root, list_file, cache, preprocess=None, transform=None):
        '''
        Args:
          root: (str) ditectory to images.
          list_file: (str/[str]) path to index file.
          cache: (DatasetCache) shared object cache
          preprocess: (function) preprocessing function before caching.
          transform: (function) image transforms.
        '''
        self.root = root
        self.cache = cache
        self.preprocess = preprocess
        self.transform = transform
        self.fnames = []
        self.labels = []

        if isinstance(list_file, list):
            # Cat multiple list files together.
            # This is especially useful for voc07/voc12 combination.
            tmp_file = '/tmp/listfile.txt'
            os.system('cat %s > %s' % (' '.join(list_file), tmp_file))
            list_file = tmp_file

        with open(list_file) as f:
            lines = f.readlines()
            self.num_imgs = len(lines)

        for line in lines:
            splited = line.strip().split()
            self.fnames.append(splited[0])
            self.labels.append(np.array([int(splited[1])]))

    def reset_memory(self):
        self.cache.reset()

    def __len__ (self):
        return self.num_imgs

    def __getitem__ (self, idx):
        '''Load image.
        Args:
          idx: (int) image index.
        Returns:
          img: (tensor) image tensor.
          labels: (tensor) class label targets.
        '''
        if self.cache.is_cached(idx):
            img, labels = self.cache.get(idx)
        else:
            # Load image and boxes.
            fname = self.fnames[idx]
            img = None
            with Image.open(os.path.join(self.root, fname)) as i:
                if i.mode != 'RGB':
                    i = i.convert('RGB')
                img = i
            if self.preprocess:
                img = self.preprocess(img)
            labels = self.labels[idx]
            labels = torch.from_numpy(labels)
            self.cache.cache(idx, img, labels)
        if self.transform:
            img = self.transform(img)
        return img, labels

# usage
manager = Manager()
cache = DatasetCache(manager)
ds = CachedListDataset(root='/path-to-data', 
                       list_file='/path-to-list-file', 
                       cache=cache,
                       preprocess=transforms.Compose([
                           transforms.Resize(256)
                       ]), 
                       transform=transforms.Compose([
                           transforms.RandomResizedCrop((224, 224), scale=[0.7, 1.0]),
                           transforms.RandomHorizontalFlip(),
                           transforms.ColorJitter(0.25, 0.25, 0.25),
                           transforms.RandomRotation(2),
                           transforms.ToTensor(),
                           transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # in case of imagenet statistics
                       ]))
dl = Dataloader(ds, ...)

```

---

<div class="post-metadata">

**Author:** ![lalithnag](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lalithnag/32/17835_2.png) [@lalithnag](https://discuss.pytorch.org/u/lalithnag)\
**Post date:** [June 22, 2021, 8:36am UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/12 "2021-06-22T08:36:45Z")

</div>

Hi thank you for this thread! Could you please let me know why creating the dataset in **init** is not a good idea?

---

<div class="post-metadata">

**Author:** ![piotr](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/piotr/32/3727_2.png) [@piotr](https://discuss.pytorch.org/u/piotr)\
**Post date:** [December 22, 2022, 12:16pm UTC](https://discuss.pytorch.org/t/best-practice-to-cache-the-entire-dataset-during-first-epoch/19608/13 "2022-12-22T12:16:10Z")

</div>

Not sure what you referring to, but init usually is expected to be fast and shouldn’t load too much data, as it can create bugs in case of multiprocessing. see [How to share data among DataLoader processes to save memory](https://discuss.pytorch.org/t/how-to-share-data-among-dataloader-processes-to-save-memory/108772)