<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">I am trying to have a generator load objects in the background, as in the following example. It hangs when trying to call <code class="notranslate">torch.zeros</code> in <code class="notranslate">split_loader_creator</code>, but if I remove the seemingly irrelevant line <code class="notranslate">torch.zeros(152*4, 168*4).float()</code> near the end, it seemingly can make progress. It also seems fine if I change <code class="notranslate">152*4</code> and <code class="notranslate">168*4</code> to much smaller numbers. This is on PyTorch 1.5.1, and I do not encounter the issue on 1.4.0. Am I somehow doing this multiprocessing incorrectly?</p>
<h2 dir="auto">To Reproduce</h2>
<p dir="auto">Run the following code:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch
import multiprocessing
import atexit

def split_loader_creator():
    for i in range(20):
        yield torch.zeros(10, 170, 70)

def background_generator_helper(gen_creator):
    def _bg_gen(gen_creator, conn):
        gen = gen_creator()
        while conn.recv():
            try:
                conn.send(next(gen))
            except StopIteration:
                conn.send(StopIteration)
                return
            except Exception:
                import traceback
                traceback.print_exc()

    parent_conn, child_conn = multiprocessing.Pipe()
    p = multiprocessing.Process(target=_bg_gen, args=(gen_creator, child_conn))
    p.start()
    atexit.register(p.terminate)

    parent_conn.send(True)
    while True:
        parent_conn.send(True)
        x = parent_conn.recv()
        if x is StopIteration:
            return
        else:
            yield x

def background_generator(gen_creator): # get several processes in the background fetching batches in parallel to keep up with gpu
    generator = background_generator_helper(gen_creator)
    while True:
        batch = next(generator)
        if batch is StopIteration:
            return
        yield batch

torch.zeros(152*4, 168*4).float()

data_loader = background_generator(split_loader_creator)
for i, batch in enumerate(data_loader):
    print(i)"><pre class="notranslate"><code class="notranslate">import torch
import multiprocessing
import atexit

def split_loader_creator():
    for i in range(20):
        yield torch.zeros(10, 170, 70)

def background_generator_helper(gen_creator):
    def _bg_gen(gen_creator, conn):
        gen = gen_creator()
        while conn.recv():
            try:
                conn.send(next(gen))
            except StopIteration:
                conn.send(StopIteration)
                return
            except Exception:
                import traceback
                traceback.print_exc()

    parent_conn, child_conn = multiprocessing.Pipe()
    p = multiprocessing.Process(target=_bg_gen, args=(gen_creator, child_conn))
    p.start()
    atexit.register(p.terminate)

    parent_conn.send(True)
    while True:
        parent_conn.send(True)
        x = parent_conn.recv()
        if x is StopIteration:
            return
        else:
            yield x

def background_generator(gen_creator): # get several processes in the background fetching batches in parallel to keep up with gpu
    generator = background_generator_helper(gen_creator)
    while True:
        batch = next(generator)
        if batch is StopIteration:
            return
        yield batch

torch.zeros(152*4, 168*4).float()

data_loader = background_generator(split_loader_creator)
for i, batch in enumerate(data_loader):
    print(i)
</code></pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">I expect this script to print the first few integers, but it just hangs at <code class="notranslate">torch.zeros</code> in <code class="notranslate">split_loader_creator</code>.</p>
<h2 dir="auto">Environment</h2>
<ul dir="auto">
<li>PyTorch Version (e.g., 1.0): 1.5.1</li>
<li>OS (e.g., Linux): Linux</li>
<li>How you installed PyTorch (<code class="notranslate">conda</code>, <code class="notranslate">pip</code>, source): conda</li>
<li>Build command you used (if compiling from source): n/a</li>
<li>Python version: 3.6.10</li>
<li>CUDA/cuDNN version: 10.2 (though I think I got the issue on 10.1 as well)</li>
<li>GPU models and configuration: Quadro RTX 6000 though the same bug happened for me on other GPUs as well</li>
<li>Any other relevant information: I suspect it's related to PyTorch, since as mentioned before I didn't encounter the issue on PyTorch 1.4.0. The issue also occurs on 1.5.0.</li>
</ul>
<h2 dir="auto">Additional context</h2>
<p dir="auto">Someone else also confirmed they could reproduce the issue: <a href="https://discuss.pytorch.org/t/pytorch-hangs-in-thread-after-large-torch-zeros-call-in-main-process/88484/3" rel="nofollow">https://discuss.pytorch.org/t/pytorch-hangs-in-thread-after-large-torch-zeros-call-in-main-process/88484/3</a></p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/ezyang/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/ezyang">@ezyang</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gchanan/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gchanan">@gchanan</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zou3519/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zou3519">@zou3519</a></p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug</h2>
<p dir="auto">The following program never terminates.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import torch
import torch.multiprocessing as mp


def foo():
    x = torch.ones((2, 50, 10))
    return torch.einsum('ijl,ikl-&gt;ijk', x, x)


if __name__ == '__main__':
    foo()
    p = mp.Process(target=foo)
    p.start()
    p.join()"><pre class="notranslate"><code class="notranslate">import torch
import torch.multiprocessing as mp


def foo():
    x = torch.ones((2, 50, 10))
    return torch.einsum('ijl,ikl-&gt;ijk', x, x)


if __name__ == '__main__':
    foo()
    p = mp.Process(target=foo)
    p.start()
    p.join()
</code></pre></div>
<p dir="auto">The behavior persists if one changes the <code class="notranslate">einsum</code> inside <code class="notranslate">foo</code> with an equivalent operation (e.g., <code class="notranslate">bmm(y, y.transpose(1,2))</code>, or <code class="notranslate">(y.unsqueeze(2) * y.unsqueeze(1)).sum(3)</code>.</p>
<p dir="auto">It doesn't reproduce, however, if one doesn't call <code class="notranslate">foo</code> inside the main block.</p>
<h2 dir="auto">Environment</h2>
<p dir="auto">PyTorch version: 1.0.1.post2<br>
Is debug build: No<br>
CUDA used to build PyTorch: 9.0.176</p>
<p dir="auto">OS: Arch Linux<br>
GCC version: (GCC) 8.2.1 20181127<br>
CMake version: version 3.13.4</p>
<p dir="auto">Python version: 3.7<br>
Is CUDA available: No<br>
CUDA runtime version: No CUDA<br>
GPU models and configuration: No CUDA<br>
Nvidia driver version: No CUDA<br>
cuDNN version: No CUDA</p>
<h2 dir="auto">Additional context</h2>
<p dir="auto">Perhaps related to <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="246557874" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/2245" data-hovercard-type="issue" data-hovercard-url="/pytorch/pytorch/issues/2245/hovercard" href="https://github.com/pytorch/pytorch/issues/2245">#2245</a>.</p>