(Multiprocessing) process hangs while looping through dataloader

Coprosmo: <p>I’ve been trying to set up parallelisation for an object detection model I’ve trained, in order to improve the throughput of the model when running on CPU. To do this, I’m roughly following <a href="https://towardsdatascience.com/this-is-hogwild-7cc80cd9b944" rel="noopener nofollow ugc">this blog post</a> on implementing Hogwild in PT.</p>
<p>Unfortunately, when running my script, the processes appear to hang while trying to iterate through the DataLoader. Iterating through the DataLoader before calling <code>mp.Process</code> works as expected, but iterating within the process causes the program to freeze.</p>
<p>I’ve provided a minimal example below:</p>
<pre><code class="lang-auto">import torch
import torch.multiprocessing as mp
from torch.utils.data import TensorDataset, DataLoader
from torch.utils.data.distributed import DistributedSampler


def loop_thru(dataloader):
    for i, _ in enumerate(dataloader):
        print(i)


if __name__ == "__main__":
    # Settings
    num_processes = 1

    # Create data
    data = torch.ones((1024,3,352,352)).float() / 255
    dataset = TensorDataset(data)

    # Start processes to loop through data
    processes = []
    for rank in range(num_processes):
        sampler = DistributedSampler(
            dataset=dataset, num_replicas=num_processes, rank=rank, shuffle=False
        )
        dataloader = DataLoader(
            dataset=dataset,
            sampler=sampler,
            batch_size=16,
        )

        p = mp.Process(target=loop_thru, args=(dataloader,))
        p.start()
        processes.append(p)

    for p in processes:
        print(f"Joining: {p}")
        p.join()
</code></pre>
<p>When I run this script, the outputs are as follows before the program hangs:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/6/4610217455cd016dae89765df71379e2420c6e97.png" data-download-href="https://discuss.pytorch.org/uploads/default/4610217455cd016dae89765df71379e2420c6e97" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/6/4610217455cd016dae89765df71379e2420c6e97_2_690x72.png" alt="image" data-base62-sha1="9ZNYVIuj6MQU3HW2gpmIW60ftIP" width="690" height="72" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/6/4610217455cd016dae89765df71379e2420c6e97_2_690x72.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/6/4610217455cd016dae89765df71379e2420c6e97_2_1035x108.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/4/6/4610217455cd016dae89765df71379e2420c6e97.png 2x" data-dominant-color="101010"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1100×116 29.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Any idea what I’m doing wrong here?</p>
<p><strong>Edit:</strong> Running pytorch=1.10.1</p>

djlindsay: <p>How long are you waiting? It ran just fine for me… produced 0-63 on separate lines.</p>

Coprosmo: <p>Just tried it again, 5 minutes and still nothing.</p>
<p>May be worth mentioning that this is running within a Docker container, though I’m not sure that should affect anything.</p>

Coprosmo: <p><strong>Update:</strong> Changed num_workers to 1, and the dataloader now iterates correctly. The process still hangs when running the data through an actual model, but the original problem is resolved.</p>

yahaha: <p>I also encountered this problem. Running on jupyter in VSCode, when num_worker is specified, the jupyter kernel will crash at the end of the DataLoader loop.<br>
I solved the problem temporarily by not setting num_worker, but I still prefer to be able to load data in multiple threads</p>

ejguan: <p>You can place you <code>Dataset</code> and <code>DataLoader</code> creation code into your <code>loop_fn</code></p>