Creating a mask tensor from an index tensor

mateja: <p>I’m trying to create a mask based on an index tensor.</p>
<p>The mask size is [6, 1, 25]<br>
The index size is [6, 1, 12]</p>
<p>First I have an index tensor <code>indices</code>:</p>
<pre><code class="lang-auto">print(indices)
tensor([[[ 0,  1,  2,  5,  6,  7, 12, 17, 18, 22, 23, 21]],

        [[ 2,  3,  4,  7,  8,  9, 15, 16, 20, 21, 22, 13]],

        [[ 0,  1,  5,  6, 10, 11, 15, 16, 17, 20, 21, 12]],

        [[ 1, 10, 15, 16, 17, 18, 20, 21, 22, 23, 24,  2]],

        [[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 13]],

        [[ 3,  4,  8,  9, 13, 14, 18, 19, 22, 23, 24, 17]]], device='cuda:0')
</code></pre>
<p>Then I allocate a mask tensor</p>
<pre><code class="lang-auto">mask = torch.cuda.ByteTensor(6, 1, 25).zero_()
</code></pre>
<p>Now I want to set only [0, 0, 0], [0, 0, 1], …, [0, 0, 23], [0, 0, 25], etc. in tensor <code>mask</code> to <code>1</code></p>
<p>I have tried:</p>
<p><code>mask[indices] = 1</code>,</p>
<p>I have also tried adapting solutions given in</p><aside class="quote quote-modified" data-post="1" data-topic="1973">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/n/8baadc/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/generating-a-meshgrid-or-itertools-product-in-pytorch/1973">Generating a meshgrid or itertools.product in pytorch</a> 
  </div>
  <blockquote>
    Hi all, 
I’m trying to generate a meshgrid directly in pytorch… 
The result I’m looking for, given two arrays of [0,1], is this: 
[[0,0],
 [0,1],
 [1,0],
 [1,1]]

You can do this with np.meshgrid: 
    grid = np.meshgrid(range(2), range(2), indexing='ij')
    grid = np.stack(grid, axis=-1)
    grid = grid.reshape(-1, 2)

and with itertools.product : 
grid = list(itertools.product(range(2),range(2)))

But does anyone reckon how you’d do that directly in pytorch? 
Thanks! 
EDIT: Hm, this seems to …
  </blockquote>
</aside>
<aside class="quote quote-modified" data-post="4" data-topic="30051">
  <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/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-select-values-using-2-index-tensor-from-a-3d-tensor/30051/4">How to select values using 2 index tensor from a 3D tensor?</a> 
  </div>
  <blockquote>
    In that case, you can just index t: 
out = t[x[:, None], y]
print(out)
&gt; tensor([[  1,   3,   5,   7,   9,  11,  13,  15],
          [ 33,  35,  37,  39,  41,  43,  45,  47],
          [ 65,  67,  69,  71,  73,  75,  77,  79],
          [ 97,  99, 101, 103, 105, 107, 109, 111],
          [129, 131, 133, 135, 137, 139, 141, 143],
          [161, 163, 165, 167, 169, 171, 173, 175],
          [193, 195, 197, 199, 201, 203, 205, 207],
          [225, 227, 229, 231, 233, 235, 237, 239]])

The last ro…
  </blockquote>
</aside>
<p>
to this problem unsuccessfully.</p>
<p>Any help would be appreciated, thank you.</p>

ptrblck: <p>I think you could use <code>scatter_</code>:</p>
<pre><code class="lang-python">mask = torch.zeros(6, 1, 25)
mask.scatter_(2, indices, 1.)
</code></pre>

mateja: <p>Wowowow that works, thank you so much <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>

Tuan-Anh_Nguyen: <p>Hi,</p>
<p>How can I do the reverse operation, meaning I want to convert a mask into indices.<br>
For example: [0, 1, 0, 1] -&gt; [1, 3]<br>
Thank you,</p>

ptrblck: <p><code>.nonzero()</code> should do the trick.</p>

sneezygiraffe: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>How do I convert a 2d mask to index  while retaining the dimensions.<br>
like if my mask is of shape Nx1000. How do i use this to create a index tensor of size (N,) as each sample in the batch may not have the same number of masked elements.</p>
<p>Thanks <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

ptrblck: <p>You won’t be able to retain the same shape or dimensions.<br>
<code>nonzero</code> will return a 2-dimensional tensor where <code>dim1</code> contains the indices for each dimension or alternatively you can return these indices as a <code>tuple</code>.</p>
<p>How should the result tensor look for your use case?<br>
E.g. assuming that not all rows in <code>N</code> contain a nonzero value, what should be stored at this particular index in <code>dim0</code>?</p>

sneezygiraffe: <p>So first let me describe my use case. I have a tensor of shape NxCxseq_len. I have an index tensor of shape NxK (basically K from each batch) where K denotes the indices from the seq_len axis and I need to pick these specific samples. I am currently using a loop to retrieve a sample corresponding to every column in the K axis but I feel there must be an easier way to go about this. Btw the indices are not the same for each sample in the batch.</p>
<p>So my resultant tensor should be of the shape NxCxK</p>

ptrblck: <p>It seems you want index a tensor of <code>[N, C, seq_len]</code> with another tensor of <code>[N, K]</code>, so I’m unsure where a mask is coming from.<br>
If I understand the use case correctly, this should work:</p>
<pre><code class="lang-python">N, C, seq_len, K = 2, 3, 4, 5

x = torch.randn(N, C, seq_len)
idx = torch.randint(0, seq_len, (N, K))
result = x[torch.arange(N)[:, None, None], torch.arange(C)[None, :, None], idx.unsqueeze(1)]
print(result.shape)
&gt; torch.Size([2, 3, 5])
</code></pre>

lolo: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Is it possible to use scatter_ with some probability? I want to do the same but instead of masking all indices from the indices matrix, I want just to mask indices with probability 15%.</p>
<p>If it can’t be done using scatter, how could I do it?</p>
<p>Thanks!</p>

ptrblck: <p><code>scatter_</code> doesn’t directly take a probability, but you could probably create the indices by sampling with a specific probability. Would this work for you?</p>