I would want to cache data in a torch.utils.data.Dataset. The simple solution is to just persist certain tensors in a member of the dataset. However, since the torch.utils.data.DataLoader class spawns multiple processes, the cache would only be local to each instance and would cause me to possibly cache multiple copies of the same tensors. Is there a way to use Python's multiprocessing library to share data between the different loader processes?
Pytorch dataset and shared memory?
Viewed 3098
1 Answers
The answer depends on your OS and settings. If you are using Linux with the default process start method, you don't have to worry about duplicates or process communication, because worker processes share memory! This is efficiently implemented as Inter Process Communication (IPC) through shared memory (some more details here). For Windows, things are more complicated. From the documentation:
Since workers rely on Python multiprocessing, worker launch behavior is different on Windows compared to Unix.
On Unix, fork() is the default multiprocessing start method. Using fork(), child workers typically can access the dataset and Python argument functions directly through the cloned address space.
On Windows, spawn() is the default multiprocessing start method. Using spawn(), another interpreter is launched which runs your main script, followed by the internal worker function that receives the dataset, collate_fn and other arguments through pickle serialization.
This means that your dynamically cached Dataset members would be automatically shared between all processes on Linux. That's great! However, on Windows, processes will not have received copies of them (they only received the Dataset upon spawning), so you should use a process communication scheme, e.g. through multiprocessing Pipe, Queue or Manager (preferred for broadcasting to multiple processes, but you would have to convert tensors to lists). This is not very efficient, besides rather bothersome to implement.
Nevertheless, there is another method: memory mapping (memmaping). This means that your objects will be written to virtual memory, and again all processes will have access to it, while a respective "shadow copy" of these objects will at some point be flushed and exist on your hard drive (can be placed in a /tmp directory). You can use memmaping with the mmap module, in which case your objects will have to be serialized as a binary file, or you can use numpy.memmap. You can find more details here.
Related