Speed up applying the same function to different tensors in Pytorch

username: hoangcuong2011
cooked: <p>I have three tensors with very different shape on the last dimension, let us say <code>tensor_a: (5, 35), tensor_b: (5, 70) and tensor_c: (5, 10)</code>. I need to apply exactly the same transformation function f() to these tensors, <code>f(tensor_a), f(tensor_b), f(tensor_c)</code> . Function f() consists of a bunch of layers. This is implemented easily with Pytorch.</p>
<p>But I noticed Gpu usage by doing sequential like that is low (around <code>50%</code>). I sped up by adding padding for <code>tensor_a</code> and <code>tensor_c</code> to make them to the same shape of <code>tensor_b, i.e.(5, 70)</code> first, then combining these three together into a single one and calling <code>f(combing_tensor)</code> instead. After that I simply slice the tensor again (:5, 5:10, 10:) to get thee resulting tensors I want. With this I noticed GPU usage is around <code>90% </code> and thus faster. However, it has a downside that it requires a lot of extra GPU memory because of padding, so it is not a perfect solution.</p>
<p>I wonder how to do it better in my case? I think need sort of embarrassing parallel in Pytorch but googling does not give me a good answer to my problem. Please let me know if there is a way to improve. Many thanks!</p>

username: Imahn
cooked: <p>How about sth like this:</p>
<pre><code class="lang-python">concatenatted_tensor = torch.cat((a, b, c), dim=1)
</code></pre>
<p>This will result in a tensor of shape <code>torch.Size([5, 115])</code>. After computing <code>f</code> on this tensor, just do</p>
<pre><code class="lang-python">a_new = concatenatted_tensor[:, :35]
b_new = concatenatted_tensor[:, 35:105]
c_new = concatentatted_tensor[:, 105:]
</code></pre>

username: hoangcuong2011
cooked: <p>Thanks. This way it does not work actually, because I need the transformation function processes a, b and c <code>separately</code>. Modifying f to take the <code>concatenatted_tensor</code> as input while keeping its subparts <code>a</code>, <code>b</code>, <code>c</code> separately is quite complicate to do.</p>

username: Imahn
cooked: <aside class="quote no-group" data-username="hoangcuong2011" data-post="1" data-topic="141601">
<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/hoangcuong2011/48/27131_2.png" class="avatar"> hoangcuong2011:</div>
<blockquote>
<p>I sped up by adding padding for <code>tensor_a</code> and <code>tensor_c</code> to make them to the same shape of <code>tensor_b, i.e.(5, 70)</code> first, then combining these three together into a single one and calling <code>f(combing_tensor)</code> instead.</p>
</blockquote>
</aside>
<p>How had you done this here? You wrote that you had padded your <code>a</code> into the shape <code>torch.Size([5, 70])</code>, how did you afterwards combine <code>a</code> and <code>b</code>? Did you do it with <code>torch.cat()</code> or <code>torch.stack()</code>?</p>

username: hoangcuong2011
cooked: <p>Let me assume two new tensors <code>a_new</code>, <code>c_new</code> have shape: <code>(5, 70)</code> after padding. So it is as simple as having a <code>combine_tensor = torch.cat([a_new, b, c_new], 0)</code>. With that combine_tensor has shape of <code>(15, 70)</code> and Function <code>f</code> takes combine_tensor as input.</p>

username: Imahn
cooked: <p>OK, I understand your problem now. I don’t see a solution to your problem… There is a Python module called <a href="https://awkward-array.org/what-is-awkward.html" rel="noopener nofollow ugc">akward</a> that is better able to deal with non-rectangular data than NumPy, but I’d say in your case, since you need to calculate <code>f(x)</code>, where <code>f</code> is a NN, you cannot use this module here.</p>
<p>Therefore, I’d go with the sequential approach, even if that has a lower utilization of the GPU as a consequence. I’d only switch to the padding option if that’s faster (you can just time both approaches).</p>