Dimension problem torch.matmul

username: lukazso
cooked: <p>Hey guys,</p>
<p>I am currently working on converting a Tensorflow project to PyTorch. At the end of the prediction, I get a translation vector field and rotation vector field of the sizes [B, 3, h, w] and [B, 3, 3, h, w] respectively.<br>
In the Tensorflow version, the outputs are [B, h, w, 3] for the translation field and [B, h, w, 3, 3] for the rotation field.</p>
<p>I now need to perform a matrix multiplication on two translation fields. The Tensorflow code would look like this:</p>
<pre><code class="lang-auto"># with rot_mat1 and rot_mat being of size **[B, h, w, 3, 3]**
r2r1 = tf.matmul(rot_mat2, rot_mat1)      
</code></pre>
<p>All working fine.</p>
<p>When I do the same without thinking in PyTorch:</p>
<pre><code class="lang-auto"># with rot_mat1 and rot_mat being of size **[B, 3, 3, h, w]**
r2r1 = tf.matmul(rot_mat2, rot_mat1)  
</code></pre>
<p>I get the error:</p>
<pre><code class="lang-auto">File "/home/student/lukas/pycharm-sync/depth-from-video/utils.py", line 131, in combine
    r2r1 = torch.matmul(rot_mat2, rot_mat1)
RuntimeError: invalid argument 6: wrong matrix size at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:558
</code></pre>
<p>I know the problem is the different dimensions in PyTorch. But do you guys have any quick fix for how I get the desired result as in the Tensorflow version? I know the answer is most likely trivial, but I am somehow stuck at the moment.</p>
<p>Thanks so much!</p>

username: ptrblck
cooked: <p>I’m not sure how TF broadcasts the matrices internally, but I guess it’s creating a new dim4 for the first tensor?<br>
If that’s the case, the second approach should work:</p>
<pre><code class="lang-python">B, H, W = 2, 5, 5

a = torch.randn(B, 3, H, W)
b = torch.randn(B, 3, 3, H, W)

out1 = torch.matmul(a.unsqueeze(1), b)
out2 = torch.matmul(a.unsqueeze(2), b)
</code></pre>
<p>Could you store the results from the TF operation and compare it to the PyTorch result?</p>

username: lukazso
cooked: <p>Thanks for your quick reply!</p>
<p>Your solution works! The problem in my case is just, that H &amp; W are of different sizes. I think i did not explain my problem properly, I am sorry. My two tensors <code>a</code> and <code>b</code> are the same size:</p>
<pre><code class="lang-auto">B, H, W = 4, 128, 416

a = torch.randn(B, 3, 3, H, W)
b = torch.randn(B, 3, 3, H, W)

</code></pre>
<p>I now tried a different approach, by transposing the tensors before the multiplication to shift them to the sizes <code>[B, H, W, 3, 3)</code>:</p>
<pre><code class="lang-auto">b = torch.transpose(b, 1, 3)
b = torch.transpose(b, 2, 4)
b.shape
&gt;&gt;&gt; torch.Size([4, 128, 416, 3, 3])

a = torch.transpose(a, 1, 3)
a = torch.transpose(a, 2, 4)
a.shape
&gt;&gt;&gt; torch.Size([4, 128, 416, 3, 3])

c = torch.matmul(a, torch.transpose(b, 3, 4))
c.shape
&gt;&gt;&gt;torch.Size([4, 128, 416, 3, 3])

</code></pre>
<p>… and then transposing c back to [B, 3, 3, H, W]. But I am not sure if this somehow messes with the whole multiplication?</p>

username: ptrblck
cooked: <p>Ah OK, in that case I misunderstood the use case.<br>
Yes, you could <code>permute</code> the inputs and reshape it back.<br>
However, from the TF docs it’s unclear to me how the 5 dimensions are handled, so could you explain the underlying operations or at least compare the outputs before continuing the porting of the code. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>