Fast way to use `map` in PyTorch?

kaiseryet: <p>So I am using PyTorch for some numerical calculation, and my problem can’t be vectorized because NestedTensor has yet to function in stable PyTorch release… Currently, I am using <code>map</code> function to do some tensor calculation. Here are two questions:</p>
<ol>
<li>Is there a more efficient way to do the parallel computation in PyTorch? e.g., I know there is a <code>tf.map_fn</code> in TensorFlow, is there anything similar in PyTorch?</li>
<li>Should I decide to use CUDA for computation, is the usage of <code>map</code> function gonna slow my algorithm?</li>
</ol>
<p>Thanks.</p>

albanD: <p>Hi,</p>
<p>I’m afraid there is no map in pytorch.<br>
If all the operations are very small, single threaded CPU will be the fastest I’m afraid.</p>
<p>If you can share your problem, maybe we can help you achieve some parallelization using the existing functions though.</p>

kaiseryet: <p>I thought <code>map</code> runs in parallel? So Python <code>map</code> doesn’t really work on PyTorch CUDA end?<br>
It’s indeed not feasible to run my problem using existing function if NestedTensor is not available… The issue is that I have to make a list of tensors of different sizes but the same dimensions, this makes <code>map</code> the only possible solution.</p>

albanD: <p>Depending on the op you need to do, you could pack then into a single Tensor with enough info to do your function <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<blockquote>
<p>I thought <code>map</code> runs in parallel?</p>
</blockquote>
<p>This is not enough to get good performance if the function needs to be applied to each entry I’m afraid.</p>
<blockquote>
<p>So Python <code>map</code> doesn’t really work on PyTorch CUDA end?</p>
</blockquote>
<p>You can still use it, it’s just that it will call the python function for every element. Which is going to not be optimal <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>

kaiseryet: <p>call function for each element… omg that’s so slow<br>
Is there any plan to maybe implement a <code>torch.map_fn</code> in the future? I mean, <code>tf.map_fn</code> definitely get some good speed from what I heard about. After NestedTensor gets implemented my problem will get solved, but create a <code>.map_fn</code> will definitely help in a broad sense.</p>

Yaroslav_Bulatov: <p>What would help here is an analogue of Jax’s <a href="https://jax.readthedocs.io/en/latest/jax.html#jax.pmap" rel="nofollow noopener">pmap</a>. I think <a class="mention" href="/u/ezyang">@ezyang</a> is experimenting with related <a href="https://github.com/pytorch/pytorch/pull/32558" rel="nofollow noopener">vmap</a></p>

albanD: <p>Well, for best perf, you want to move your function into the inner loop in cpp or cuda code.<br>
Which you cannot do using a python function.</p>
<p>Also remember that NestedTensors are built on top of pytorch. So if you can do your code with NestedTensor, you can do it without <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

kaiseryet: <aside class="quote no-group" data-username="albanD" data-post="7" data-topic="70814">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>move your function into the inner loop in cpp or cuda code</p>
</blockquote>
</aside>
<p>Too complicated for my purpose…</p>

albanD: <p>It is for most use case.</p>
<p>But as mentioned above, there is very little you cannot do with pytorch-only functions now. If you want to share a code sample with your function implemented with a for-loop, we might be able to help.</p>

kaiseryet: <p>An example of my purpose is, lets say we hate a list of square matrices of two dimensions, called <code>A</code>, a list of some matrices of two dimensions, <code>B</code>, of the same list length as <code>A</code>. I’m trying to write a function to find the list of quadratic form <code>B.T@A@B</code>.</p>
<p>e.g.</p>
<pre><code class="lang-auto">A = [torch.randn(i, i) for i in range(2001)]
B = [torch.randn(i, 3) for i in range(2001)]
[B[i].T@A[i]@B[i] for i in range(2001)] # this is the target!
</code></pre>

albanD: <p>And the sizes are wildly different?<br>
Does packing all the A matrices in a single Tensor with 0 padding very wasteful in your case?</p>

kaiseryet: <p>See my updated with an example, I would say padding with 0 might well be very wasteful… that’s one example, second example could be finding the inverse, which will be even worse in that situation.</p>

albanD: <p>Right but for such tasks, you want to do a very large number of small tasks.<br>
So you won’t be able to run this on the GPU as it will be super slow.<br>
And you will have to run this one by one on the CPU in a for-loop. Moving that for loop from python to cpp will speed it up but won’t give you "good’ perf I’m afraid…</p>
<p>You can use the cpp inline extensions to try this:</p>
<pre><code class="lang-python">import torch
from torch.utils import cpp_extension
import time

cpp_source = """
std::vector&lt;torch::Tensor&gt; test_fn(std::vector&lt;torch::Tensor&gt; A, std::vector&lt;torch::Tensor&gt; B) {
    std::vector&lt;torch::Tensor&gt; result(A.size());

    for (int i=0; i&lt;A.size(); ++i) {
        result[i] = A[i].mm(B[i]);
    }

    return result;
}
"""

def py_version(A, B):
    C = []
    for a, b in zip(A, B):
        C.append(a.mm(b))
    return C

mod = cpp_extension.load_inline(
    name="mod",
    cpp_sources=cpp_source,
    functions=["test_fn",],
)


for max_val in [10, 100, 1000, 1500]:
    print("")
    print("TESTING FOR {}".format(max_val))
    A = [torch.randn(i, i) for i in range(1, max_val)]
    B = [torch.randn(i, 3) for i in range(1, max_val)]

    t1 = time.time()
    res = mod.test_fn(A, B)
    cpp_time = time.time() - t1
    t1 = time.time()
    res = py_version(A, B)
    py_time = time.time() - t1

    print("cpp time: {}".format(cpp_time))
    print("py time: {}".format(py_time))
</code></pre>

kaiseryet: <p><img src="https://discuss.pytorch.org/images/emoji/apple/pensive.png?v=9" title=":pensive:" class="emoji" alt=":pensive:"> <img src="https://discuss.pytorch.org/images/emoji/apple/pensive.png?v=9" title=":pensive:" class="emoji" alt=":pensive:"> <img src="https://discuss.pytorch.org/images/emoji/apple/pensive.png?v=9" title=":pensive:" class="emoji" alt=":pensive:"> finally have to brought up cpp…</p>
<p>when will NestedTensor implement in the stable release though? <a class="mention" href="/u/cpuhrsch">@cpuhrsch</a> pls?</p>

albanD: <p>My example is more to show that you most likely don’t want cpp:</p>
<pre><code class="lang-auto">TESTING FOR 1500
cpp time: 0.2952902317047119
py time: 0.27726316452026367
</code></pre>
<p>This is the timings I get for a simple mm with the inputs you gave <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>
<p>cc <a class="mention" href="/u/cpuhrsch">@cpuhrsch</a> for the nested tensor release <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=9" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"></p>

kaiseryet: <p>sorry, scared when read cpp.</p>
<p>But wait, Isn’t <code>map</code> supposed to be faster than <code>for</code> due to parallelisation?</p>

albanD: <p>Well that depends on how large each op is. If they are big enough, we already paralellize each op.<br>
In the case where they are super super small, yes you would gain from a parallelization on top.</p>
<p>But to get back to the main question:</p>
<ul>
<li>we don’t have a map</li>
<li>I don’t think there is any plan to make one
<ul>
<li>It would be a lot of work</li>
<li>It would have limited benefit has it would only help very small, cpu-only tasks.</li>
</ul>
</li>
<li>If your ops are small, you don’t want to use CUDA</li>
</ul>

kaiseryet: <p>okay fair enough. I’m just giving an example here.</p>

hypnagogic: <p>Just a quick thought from a cursory skim of this problem. I’s wondering whether a purely <code>multiprocessing</code> flow may help (on cpus) here:</p>
<ol>
<li>chunk the lists (A, B)</li>
<li>use a <code>multiprocessing.Pool.map(func, zip(A, B))</code>
</li>
</ol>
<p>something like (quick pseudocode)</p>
<pre><code class="lang-auto">def chunks(a:List, b:List, n:int) -&gt; Generator;

nproc = os.cpu_count() - 1
with Pool(proceses=nproc) as pool:
  func = functools.partial(process_matrices, *params)
  pool.map(func, zip(a_chunk, b_chunk))

</code></pre>