Shared memory with torch.multiprocessing

username: wamreyaz
cooked: <p>Hey folks,</p>
<p>I have a server with large amounts of RAM, but slow storage and I want to speed up training by having my dataset in the RAM. I also use <code>DDP</code> which means there are going to be multiple processes per GPU. On top of that, I use multiple <code>num_workers</code> in my <code>dataloader</code> so having a simple Python list as a  caxhe would mean multiple caches which eats up a lot of memory.</p>
<p>The natural solution is to use shared memory. And this is how I use it</p>
<ol>
<li>In the launch process, do</li>
</ol>
<pre><code class="lang-auto">if __name__ == '__main__':
    import argparse
    import os
    import torch.multiprocessing as mp
    import ctypes

    shared_base = mp.Array(ctypes.c_byte, 80000*3*256*256, lock=True)
    with shared_base.get_lock():
        shared_array = np.ctypeslib.as_array(shared_base.get_obj())
        img_cache = shared_array.reshape(80000, 256, 256, 3)

    use_cache = mp.Array(ctypes.c_float, 1, lock=False)
    use_cache[0] = -3.14
</code></pre>
<ol start="2">
<li>This cache is sent to each process as</li>
</ol>
<pre><code class="lang-auto">mp.spawn(main, nprocs=ngpus_per_node, args=(args, img_cache, use_cache))
</code></pre>
<ol start="3">
<li>Each process takes it this shared memory and gives it to a dataset object</li>
</ol>
<pre><code class="lang-auto">        dset = SVAE_FFHQ(args.data_folder, transform, 32, 64, args.hidden_size, img_cache, use_cache)

</code></pre>
<ol start="4">
<li>The <code>SVAE_FFHQ</code> class does looks like this:</li>
</ol>
<pre><code class="lang-auto">class SVAE_FFHQ(data.Dataset):
    def __init__(self, root_dir, transform=None, top_size=32, bottom_size=64, dim=256, img_cache=None, use_cache=None):
        super().__init__()
        ...
        self.img_cache = img_cache
        self.use_cache = use_cache


    def _use_cache(self):
        self.use_cache[0] = 3.14
        print('Using cache')

    def __getitem__(self, idx):
        path, lbl = self.dset.samples[idx]

        if self.use_cache[0] &lt; 0:
            with open(path, 'rb') as f:
                    img = Image.open(f)
                    img = img.convert('RGB')
                    img = img.resize((256, 256), Image.LANCZOS)
            
            self.img_cache[idx] = deepcopy(np.asarray(img))
            del img
        
        return self.transform(Image.fromarray(self.img_cache[idx], 'RGB'))
</code></pre>
<p>This to me seems fine, but what happens is</p>
<ol>
<li>The shared memory is pickled and not replicated across the multiple spawned processes which means my memory requirments increase with the number of processes spawned.</li>
<li>This isn’t any faster than reading data off of slow HDDs</li>
</ol>
<p>Any insight into these problems?</p>
<p>Thanks!</p>

username: iffiX
cooked: <ol>
<li>I cannot replicate your first problem with the below code snippet, memory is not pickled as you can see in the screenshot, only the main process holds the 1GB shared memory array:</li>
</ol>
<pre><code class="lang-auto">import ctypes
import time
import numpy as np
import multiprocessing as mp


def subproc(array):
    with array.get_lock():
        np_array = np.ctypeslib.as_array(array.get_obj())
        print(np_array[1000])

    # keep process showing in "top"
    begin = time.time()
    while time.time() - begin &lt; 10:
        a = 1 * 100


if __name__ == "__main__":
    array = mp.Array(ctypes.c_byte, 1000*1024*1024, lock=True)
    with array.get_lock():
        np_array = np.ctypeslib.as_array(array.get_obj())
        np_array.fill(100)
    print("allocated")
    p = mp.Process(target=subproc, args=(array,))
    p2 = mp.Process(target=subproc, args=(array,))
    p.start()
    p2.start()
    print("started")

    # keep process showing in "top"
    begin = time.time()
    while time.time() - begin &lt; 10:
        a = 1 * 100
    p.join()
    p2.join()
    print("joined")
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/b/1b3bcfa32a454d99f05015ae09fb325ab2e2d59c.png" alt="image" data-base62-sha1="3SV2fKZvFJc6JvEJKYvFyIesbsE" width="536" height="68"></p>
<ol start="2">
<li>I think the second question might be related to <code>Image.fromarray</code>, see this <a href="https://stackoverflow.com/questions/11172929/can-you-create-a-pil-image-with-cython-without-copying-memory-redundantly" rel="noopener nofollow ugc">issue</a></li>
</ol>

username: wamreyaz
cooked: <p>Thanks for implementing this by yourself.</p>
<p>I must ask a couple of questions though:</p>
<ol>
<li>What was the version of Python you used? I know that some stuff changed in how Python pickles data since version 3.8.</li>
<li>I use mp.spawn to start the processes, where mp is imported from torch.multiprocessing.</li>
</ol>
<p>Could that be a problem?</p>
<p>Anyways this looks interesting. Thanks again.</p>

username: iffiX
cooked: <p>I am using the default python3.5.2 installation from ubuntu16.04</p>
<p>I have added <code>mp.set_start_method("spawn")</code> and still cannot reproduce your issue, it would be better if you can share a minimal problematic code snippet <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=9" title=":blush:" class="emoji" alt=":blush:"></p>

username: wamreyaz
cooked: <p>Thanks again.</p>
<p>Let me also try your code on my machine. And create a minimal code implementation that I can share here.</p>

username: wamreyaz
cooked: <p>This is the code I run</p>
<pre><code class="lang-auto">import ctypes
import time
import numpy as np
import torch.multiprocessing as mp

def subproc2(gpu, array):
    with array.get_lock():
        np_array = np.ctypeslib.as_array(array.get_obj())
        print(np_array[1000])
        if gpu == 0:
            np_array[999] = 0
        elif gpu == 1:
            np_array[1000] = 1 

    # keep process showing in "top"
    begin = time.time()
    while time.time() - begin &lt; 10:
        a = 1 * 100

    return 0

if __name__ == "__main__":
    mp.set_start_method('spawn')
    array = mp.Array(ctypes.c_byte, 1000*1024*1024, lock=True)
    with array.get_lock():
        np_array = np.ctypeslib.as_array(array.get_obj())
        np_array.fill(100)

    print("allocated")
    mp.spawn(subproc2, args=(array,), nprocs=2)

    # keep process showing in "top"
    print(np_array[999:10001])
    begin = time.time()
    while time.time() - begin &lt; 100:
        a = 1 * 100
    
    print('done')
</code></pre>
<p>The observations are :</p>
<ol>
<li>Memory is only ever allocated in the main process.</li>
<li>If I don’t set the start method to <code>spawn</code> on Linux, I get a <code>SIGSEGV</code> even on access in <code>subproc2</code>.</li>
</ol>
<p>Now what I want to test next is if I pass this memory on to dataloaders as is - an <code>np.array</code>, if it will increase memory consumption.</p>
<p>Let me do that now.</p>

username: wamreyaz
cooked: <pre><code class="lang-auto">import ctypes
import time
import numpy as np
import torch.multiprocessing as mp
import torch
from torch.utils.data import Dataset

class SharedDataset1(Dataset):
    def __init__(self, shared_mem):
        super(SharedDataset1, self).__init__()
        self.shared_mem = shared_mem

    def __len__(self):
        return 10000

    def __getitem__(self, idx):
        return torch.randn(3, 32, 32)

def np_before(gpu, array, num_workers):
    with array.get_lock():
        np_array = np.ctypeslib.as_array(array.get_obj())

    dataset = SharedDataset1(np_array)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=10, num_workers=num_workers)

    for img in dataloader:
        c = img+0.1
        time.sleep(3)

if __name__ == "__main__":
    mp.set_start_method('spawn')
    array = mp.Array(ctypes.c_byte, 10000*1024*1024, lock=True)
    with array.get_lock():
        np_array = np.ctypeslib.as_array(array.get_obj())
        np_array.fill(100)

    print("allocated")
    mp.spawn(np_before, args=(array, 1), nprocs=2)
    print("started")

    begin = time.time()
    while time.time() - begin &lt; 100:
        a = 1 * 100
</code></pre>
<p>This gives the following error</p>
<blockquote>
<p>Traceback (most recent call last):<br>
File “mem_data.py”, line 64, in <br>
mp.spawn(np_before, args=(array, 1), nprocs=2)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/site-packages/torch/multiprocessing/spawn.py”, line 200, in spawn<br>
return start_processes(fn, args, nprocs, join, daemon, start_method=‘spawn’)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/site-packages/torch/multiprocessing/spawn.py”, line 158, in start_processes<br>
while not context.join():<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/site-packages/torch/multiprocessing/spawn.py”, line 119, in join<br>
raise Exception(msg)</p>
</blockquote>

username: wamreyaz
cooked: <blockquote>
<p>Traceback (most recent call last):<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/site-packages/torch/multiprocessing/spawn.py”, line 20, in _wrap<br>
fn(i, *args)<br>
File “/home/parawr/Projects/shared/mem_data.py”, line 40, in np_before<br>
for img in dataloader:<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/site-packages/torch/utils/data/dataloader.py”, line 279, in <strong>iter</strong><br>
return _MultiProcessingDataLoaderIter(self)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/site-packages/torch/utils/data/dataloader.py”, line 719, in <strong>init</strong><br>
w.start()<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/multiprocessing/process.py”, line 112, in start<br>
self._popen = self._Popen(self)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/multiprocessing/context.py”, line 223, in _Popen<br>
return _default_context.get_context().Process._Popen(process_obj)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/multiprocessing/context.py”, line 284, in _Popen<br>
return Popen(process_obj)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/multiprocessing/popen_spawn_posix.py”, line 32, in <strong>init</strong><br>
super().<strong>init</strong>(process_obj)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/multiprocessing/popen_fork.py”, line 20, in <strong>init</strong><br>
self._launch(process_obj)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/multiprocessing/popen_spawn_posix.py”, line 47, in _launch<br>
reduction.dump(process_obj, fp)<br>
File “/home/parawr/.conda/envs/faclab/lib/python3.7/multiprocessing/reduction.py”, line 60, in dump<br>
ForkingPickler(file, protocol).dump(obj)<br>
OverflowError: cannot serialize a bytes object larger than 4 GiB</p>
</blockquote>

username: wamreyaz
cooked: <p>This thing works</p>
<pre><code class="lang-auto">def np_after(gpu, array, num_workers):
    dataset = SharedDataset1(array)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=10, num_workers=num_workers)

    for ii, img in enumerate(dataloader):
        print(ii)
        c = img+0.1
        time.sleep(1)
    return 0

class SharedDataset2(Dataset):
    def __init__(self, shared_mem):
        super(SharedDataset2, self).__init__()
        with shared_mem.get_lock():
            self.shared_mem = np.ctypeslib.as_array(shared_mem.get_obj())

    def __len__(self):
        return 10000

    def __getitem__(self, idx):
        return torch.randn(3, 32, 32)
</code></pre>
<p>I think it might have to do with what Python pickles - the array/shared_mem is a <code>SynchronizedArray</code> object which Python <em>knows</em> no to pickle but share, while in the <code>np_before</code> function it is now an <code>np.ndarray</code> which Python tries to pickle.</p>
<p>Let me now try and fix the PIL issue. I will benchmark it to see if it actually speed up data-loading.</p>

username: wamreyaz
cooked: <p>I will also have to check validity - I wouldn’t want the workers from the dataloader overwriting each others’  data</p>

username: iffiX
cooked: <p>The problem source is correct, since the dataset passed to <code>Dataloader</code> will be distributed to workers, and therefore will be pickled. Good work, looking forward to your result. <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=9" title=":blush:" class="emoji" alt=":blush:"></p>