Is there a way to parallelize independent sequential steps?

username: pranav
cooked: <p>I am currently working with a model which has a fixed set of heavy operations applied on different versions of the same input. It would be really nice if I could perform these operations in parallel as each one of them takes a big chuck of total execution time. In Theano, with graph optimizations, this would have been automatically executed in parallel. But I am not sure how to parallelize these operations in PyTorch. I looked into Python’s multiprocessing module and PyTorch’s wrapper for it. But I am not sure if it’s usage would maintain the autograd’s graph integrity and how they’d work while performing backprop (<em>which is actually an even more expensive than forward propagation</em>).</p>
<p>An NDA prevents me from sharing the exact model details and code, but here is a representative example of the model that captures the problem.</p>
<pre><code class="lang-python">def forward(x):
    # transform the input in 4 different ways - transform is a lightweight operation
    x1 = transform(x, 0, 0)
    x2 = transform(x, 0, 1)
    x3 = transform(x, 1, 0)
    x4 = transform(x, 1, 1)

    # process each transformed input by the same function
    # my_func here is time consuming and the current bottleneck, but is internally optimized
    # ideally, all four `my_func` calls below should execute in parallel
    a = my_func(x1)
    b = my_func(x2)
    c = my_func(x3)
    d = my_func(x4)

    result = a + b + c + d
    
    return result
</code></pre>
<p>Is there any way I could speed up this function? Not being able to parallelize this hurts me even more when I stack many such steps, which I am planning to do next. So any suggestions to speed this up in any way would help a lot.</p>
<p>Thanks!</p>

username: pranav
cooked: <p>After seeing some discussion on Twitter, it looks like having lazy evaluation would defeat some of the goals of a truly dynamic tensor manipulation framework. But for cases like the example above at least, having a declarative way that would allow for parallel execution of several operations would be very useful.</p>

username: jekbradbury
cooked: <p>Depending on how complex my_func is, you may be able to get some mileage out of wrapping it in <code>nn.DataParallel</code> and running each of them on separate GPUs if you have several available. But you can also rewrite <code>my_func</code> to take a list of variables and manually batch certain suboperations together; that’s likely to be the fastest approach.</p>

username: rteja1113
cooked: <p>Hi pranav, it seems that static graphs have upper hand at these kind of parallel tasks. In the official torchvision InceptionV3, inception modules are implemented serially.</p>

username: sad_robot
cooked: <aside class="quote no-group" data-username="pranav" data-post="1" data-topic="3360">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/pranav/48/185_2.png" class="avatar"> pranav:</div>
<blockquote>
<p>In Theano, with graph optimizations, this would have been automatically executed in parallel. But I am not sure how to parallelize these operations in PyTorch.</p>
</blockquote>
</aside>
<p>7 Years later and the answer is (for GPU operations) that it will happen automatically! Execution is eager but also non-blocking (for GPU ops) so the GPU will queue each operation then immediately return control the CPU to queue further operations. If the operations are independent then the GPU can &amp; will run them in parallel automatically.</p>