Dot product batch-wise

avijit_dasgupta: <p>I have two matrices of dimension (6, 256). I would like to calculate the dot product row-wise so that the dimensions of the resulting matrix would be (6 x 1). <code>torch.dot</code> does not support batch-wise calculation. Any efficient way to do this?</p>

richard: <p>Each row is a vector with 256 elements; what do you mean by dot product? Do you just want to multiply all of those elements together?</p>

SimonW: <p><code>torch.bmm(A.view(6, 1, 256), B.view(6, 256, 1))</code> should do the trick!</p>
<p><a href="http://pytorch.org/docs/0.2.0/torch.html#torch.bmm" class="onebox" target="_blank">http://pytorch.org/docs/0.2.0/torch.html#torch.bmm</a></p>

avijit_dasgupta: <p>Yeah! That would do. There is no direct function then, right?</p>

SimonW: <p>I don’t think so. Well, this is very efficient because there is no copying to massage the data <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ecolss: <p>OK, let’s say, on mac osx, cpu,</p>
<pre><code class="lang-auto">In [166]: a = th.Tensor(np.random.rand(100000, 10))

In [167]: %timeit th.sum(a*a, dim=1).sum()
2.8 ms ± 39.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [168]: %timeit th.bmm(a.view(-1,1,10), a.view(-1,10,1)).sum()
64.2 ms ± 666 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
</code></pre>
<p><code>bmm</code> is almost 30 times slower, and furthermore <code>bmm</code> makes <code>backward</code> much more slower, what?</p>

Stone: <p>In addition, the results are similar on GPU.<br>
For GPU version, <code>bmm</code> invokes separate CUDA kernels for each matrix multiplication, which in this case, <strong>10000</strong> kernel launches. Then the function calling overhead dominates the total computation time.</p>

SimonW: <p>yeah, it might not be specifically optimized for this case. thanks for doing the benchmark.</p>

SimonW: <p>I just checked bmm is a single batched gemm kernel call. It’s not doing 10000 kernel launches.</p>

Stone: <p>Yeah you’re right, it uses a single batched kernel. Actually the CPU version is a loop in <em>batch</em> dimension.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/29849e428cba03f4a9e24a157781ce822512c39e/aten/src/TH/generic/THTensorMath.cpp#L2308" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/29849e428cba03f4a9e24a157781ce822512c39e/aten/src/TH/generic/THTensorMath.cpp#L2308" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/29849e428cba03f4a9e24a157781ce822512c39e/aten/src/TH/generic/THTensorMath.cpp#L2308</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="2298" style="counter-reset: li-counter 2297 ;">
<li>  THTensor_(resizeAs)(result, t);</li>
<li>  if (beta != 0.0) {</li>
<li>    THTensor_(copy)(result, t);</li>
<li>  }</li>
<li>}</li>
<li>
</li>
<li>THTensor *matrix1 = THTensor_(new)();</li>
<li>THTensor *matrix2 = THTensor_(new)();</li>
<li>THTensor *result_matrix = THTensor_(new)();</li>
<li>
</li>
<li class="selected">for (batch = 0; batch &lt; THTensor_(size)(batch1, 0); ++batch) {</li>
<li>  THTensor_(select)(matrix1, batch1, 0, batch);</li>
<li>  THTensor_(select)(matrix2, batch2, 0, batch);</li>
<li>  THTensor_(select)(result_matrix, result, 0, batch);</li>
<li>
</li>
<li>  THTensor_(addmm)(result_matrix, beta, result_matrix, alpha, matrix1, matrix2);</li>
<li>}</li>
<li>
</li>
<li>THTensor_(free)(matrix1);</li>
<li>THTensor_(free)(matrix2);</li>
<li>THTensor_(free)(result_matrix);</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


bananacode: <p>If anyone came across this post via a google search, I suggest they check out the following github issue:<br>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/18027" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <a href="https://github.com/sidazhang" rel="nofollow noopener">
<img src="https://avatars1.githubusercontent.com/u/3519115?v=2&amp;s=96" class="thumbnail onebox-avatar" width="60" height="60">
</a>

<h4><a href="https://github.com/pytorch/pytorch/issues/18027" target="_blank" rel="nofollow noopener">Issue: Support Batch Dot Product</a></h4>

<div class="date" style="margin-top:10px;">
	<div class="user" style="margin-top:10px;">
	opened by <a href="https://github.com/sidazhang" target="_blank" rel="nofollow noopener">sidazhang</a>
	on <a href="https://github.com/pytorch/pytorch/issues/18027" target="_blank" rel="nofollow noopener">2019-03-14</a>
	</div>
	<div class="user">
	</div>
</div>

<pre class="content" style="white-space: pre-wrap;">🚀 Feature
Support batch dot product
Motivation
Commonly used operation
Alternatives
Currently, we can do this with bmm. But I think it makes sense to just...</pre>

<div class="labels">
 	<span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">feature</span>
 	<span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Given two batches of vectors A,B, it is the fastest to just compute (A*B).sum(-1)</p>

tengerye: <p>I add another method using <code>matmul()</code> with <code>transpose()</code>. The order is from faster to slower:</p>
<pre><code class="lang-auto">a = torch.rand(2, 4)

%timeit (a*a).sum(1)
# 4.26 µs ± 21.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit torch.matmul(a, a.t()).diag()
# 6.81 µs ± 365 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit torch.bmm(a.view(2, 1, 4), a.view(2, 4, 1)).view(2, 1)
# 16.2 µs ± 156 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)</code></pre>

jshtok: <p>In general,<br>
torch.bmm(A.unsqueeze(dim=1), B.unsqueeze(dim=2)).squeeze()</p>