Data loader multiprocessing slow on macOS

username: adamjstewart
cooked: <p>I’ve noticed a significant slowdown when using <code>num_workers &gt; 0</code> on macOS for Python 3.8+. With the following benchmarking script:</p>
<pre><code class="lang-python">import sys
import time

from torch.utils.data import DataLoader, Dataset


class MyDataset(Dataset):
    def __getitem__(self, i):
        return i

    def __len__(self):
        return 5


if __name__ == '__main__':
    ds = MyDataset()
    dl = DataLoader(ds, num_workers=int(sys.argv[1]))

    start = time.time()
    for i in dl:
        continue
    end = time.time()
    print(end - start)
</code></pre>
<p>I see these times on macOS:</p>
<pre><code class="lang-auto">$ python test.py 0
0.0023140907287597656
$ python test.py 1
9.695190906524658
$ python test.py 2
14.391160249710083
$ python test.py 3
18.87164807319641
</code></pre>
<p>I would expect these to take approximately the same time. Obviously, there is some overhead when creating and offloading parallel workers, but this is significantly greater than on Linux:</p>
<pre><code class="lang-auto">$ python test.py 0
0.0008292198181152344
$ python test.py 1
0.025320768356323242
$ python test.py 2
0.03378176689147949
$ python test.py 3
0.030670166015625
</code></pre>
<p>The issue seems to be that starting with Python 3.8, the default multiprocessing <a href="https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods" rel="noopener nofollow ugc">start method</a> changed from <em>fork</em> to <em>spawn</em>. The main difference is that with <em>spawn</em>, all resources of the parent need to be pickled so they can be inherited by the child. It seems that this is where the slowdown is coming from, but I can’t figure out how to speed up pickling times. Do I need to add a <code>__reduce__</code> attribute to any of these objects or something? Or is this overhead always required and there’s nothing I can do about it? I realize that this slowdown is miniscule compared to the time involved in actually loading real data and passing it through a model, but it’s really slowing down my unit tests and affecting other benchmarking I’m trying to do.</p>

username: jia.kai
cooked: <p>I also encountered this problem. The slowest part seems to be worker startup/shutdown (you can add <code>print(i)</code> before <code>continue</code>; data retrieving is quite fast). I don’t have time to investigate the root cause. My temporary workaround is to use <code>mp.set_start_method('forkserver')</code>.</p>

username: laclouis5
cooked: <p>I also noticed that DataLoader shutdown is very slow (between 5s and 10s), even in a recent environment (MacBook Pro 14" with M1 Pro running PyTorch 2.0.0). As noted by <a class="mention" href="/u/jia.kai">@jia.kai</a>, the issue is that PyTorch multiprocessing uses the <code>spawn</code> method on macOS (and also Linux I guess).</p>
<p>The answer from <a class="mention" href="/u/jia.kai">@jia.kai</a> works fine, however, I would recommend using the dedicated <code>DataLoader</code> parameter instead:</p>
<pre><code class="lang-python">from torch.utils.data import DataLoader

def main():
  dataset = ...

  dataloader = DataLoader(
    dataset, num_workers=4, 
    multiprocessing_context="forkserver",
  )

if __name__ == "__main__":
  main()
</code></pre>
<p>If used in a training loop, I would also recommend passing <code>persistent_workers=True</code> to the <code>DataLoader</code> input paramsin order to avoid recreating the worker processes at the beginning of each dataset iteration.</p>
<p>Also, instead of the Python <code>multiprocessing</code> library I would recommend using <a href="https://pytorch.org/docs/stable/multiprocessing.html" rel="noopener nofollow ugc"><code>torch.multiprocessing</code></a> when it makes sense to do so.</p>

username: kbwiley
cooked: <p>Oh man! Thank you! Both of your suggestions made immediate improvements.</p>
<p><code>multiprocessing_context="forkserver"</code> eliminated the huge delay (5s/allocated-worker) wrapping up the final iteration over the batches of the DataLoader, and then the subsequent suggestion of <code>persistent_workers=True</code> further removed some apparently unnecessary slowdowns in all batches after the first one.</p>
<p>I wish this would be better incorporated into the source. If PyTorch detects that it is running on a Mac, it should default to the more performant <code>multiprocessing_context</code>, unless there is some good reason for it to default to a value that seemingly renders PyTorch all but usless on a Mac (I’m new to PyTorch and am willing to believe there is a justification for the default value, but I still don’t see how to get any practical utility out of PyTorch on a Mac without altering the mp context as described; perhaps I am misunderstanding something.). As for <code>persistent_workers</code>, if it is all gain at no obvious cost, I’m curious what the default value is the slower option. Again, I suspect I am not understanding it and there is good reason for the default value to be the seemingly less desirable value in all scenarios I can imagine, but again, since I’m new, perhaps I’m just unaware of reasonable reasons for the default to be as it is instead of the opposite value.</p>
<p>Thanks again!</p>