Select data through a mask

Paulo_Mann: <p>Hello,</p>
<p><strong>Question 1)</strong></p>
<p>If I have a <strong>Tensor</strong> of size <strong>(A, B, C)</strong>  and a <strong>mask</strong> of size <strong>(A, B)</strong>. The mask only has values of 1 or 0 which indicates whether we should use an example of the dimension B of my tensor.</p>
<p>In practice, I would like to select all (A, B, C) examples of my <strong>Tensor</strong> that have the respective mask values of 1 and ignore the rest. After that, I would like to take the mean of this (A, B, C) with respect to dim=1, which I can do through torch.mean(tensor, dim=1).</p>
<p>Example:</p>
<p>embeddings.size() -&gt; torch.Size([2, 3, 1024])<br>
mask.size()            -&gt; torch.Size([2, 3])<br>
mask                      -&gt; tensor([[1, 1, 0], [1, 1, 1]])</p>
<p>I would like to have for the A = 0: embeddings[0, 0:2, :] | and for A = 1: embeddings[0, :, :], and take, after that, the mean over the dim=1.</p>
<p>OBS: I know I can do it with for loops, but I want to know a cleaner and more efficient way to do this.</p>
<p><strong>Question 2)</strong></p>
<p>All this calculation goes under the <code>forward()</code> function of my model. All this slicing and mean go to the grad_fn function, is that right? Could it be influencing all the backpropagation in a wrong way? For a concrete example, calculating the mean in my application is taking the mean of word embeddings to get a sentence embedding.</p>
<p>Thanks!</p>

MariosOreo: <pre><code class="lang-auto">mask = mask.unsqueeze(-1)
mask_embeddings = embeddings * mask.float()
result = mask_embeddings.mean(dim=1)
</code></pre>
<p>Does this fit your use case?</p>
<p>AFAIK, slicing and mean will work with autograd.</p>

Paulo_Mann: <p>Hello <a class="mention" href="/u/mariosoreo">@MariosOreo</a>, first, thanks for the answer.</p>
<p>It works, but in the <code>mask_embeddings.mean(dim=1)</code> you gonna consider the vectors in the B dimension that are 0 in the mask for the mean… Although they don’t contribute for the sum, it is going to contribute as a value for the division part of the mean, right?</p>
<p>In the example I gave:</p>
<blockquote>
<p>embeddings.size() → torch.Size([2, 3, 1024])<br>
mask.size() → torch.Size([2, 3])<br>
mask → tensor([[1, 1, 0], [1, 1, 1]])</p>
</blockquote>
<p>The example 0 of the batch is going to have a size of <code>torch.Size([2, 1024])</code>. Whereas the example 1 of the batch <code>torch.Size([3, 1024])</code>, and I would like to take the mean for the example 0 over the 2 word embeddings, and over the 3 word embeddings for the example 1.</p>

MariosOreo: <p>In your use case, I am afraid that <code>torch.mean()</code> cannot fit our requirement.<br>
You could have a try on this:</p>
<pre><code class="lang-auto">sum_mask_embeddings = mask_embeddings.sum(dim=1)
for dim in range(mask.size(0)):
    average_mask_embeddings[dim] = sum_mask_embeddings[dim] / mask[dim].sum()
</code></pre>

Paulo_Mann: <p>If I don’t want this mean calculation to take part of the autograd process, can I wrap your code with <code>torch.no_grad()</code>?</p>
<pre><code class="lang-auto">with torch.no_grad():
    sum_mask_embeddings = mask_embeddings.sum(dim=1) 
    for dim in range(mask.size(0)):
        average_mask_embeddings[dim] = sum_mask_embeddings[dim] / mask[dim].sum()
</code></pre>
<p>Does this makes sense? Do you think it can have a beneficial impact on my embeddings?</p>

MariosOreo: <p>Yes, you are right. If you don’t want take this part into autograd process, you can wrap it into <code>torch.no_grad()</code> scope.</p>
<p>And we often use <code>torch.no_grad()</code> to avoid some error (i.e. modify the parameters of the network inplace). For whether it can have beneficial impacts on your embeddings, I think it depends on your use case.</p>