How to slice the tensor with the following mask in the easiest way?

jack_S: <p>I have a tensor, say representing some images, with a shape <code>[batch_size, channel, height, width]</code>, and a mask tensor with a shape <code>[batch_size, channel]</code>. All elements in the mask tensor are bools. <code>True</code> means I want to slice that channel while <code>False</code> means I don’t. For example:</p>
<pre><code class="lang-auto">image = torch.randn(100, 3, 224, 224)
mask = torch.randn(100, 3).ge(0.5)
</code></pre>
<p>I wanted to slice <code>image</code> in the easiest way so I tried <code>image[:, mask]</code> but this didn’t work.</p>
<p>I realized that I have to add an extra condition: the number of <code>True</code>s in <code>dim = 1</code> is the same. Thus, the shape of the sliced tensor is <code>[batch_size, num_of_true, height, width]</code>.</p>

ptrblck: <p>As you’ve already described, directly indexing <code>image</code> won’t work as the output tensor would have a variable shape in <code>dim1</code>. You could use <code>image[mask]</code> which will output <code>[num_of_true, height, width]</code>, but unsure if you could work with this result.</p>

jack_S: <p>Now I have to use <code>for</code> to do an element-wise work, but I’m not sure if  using <code>for</code> slows the speed.</p>

J_Johnson: <p>Pytorch can’t have tensors which are varied size on a given dim. Here are some ways you can approach this, depending on your desired outcome:</p>
<pre><code class="lang-auto">import torch

#apply the filter evenly on all batches the same
images=torch.randn(100, 3, 224, 224)
mask = torch.randn(3).ge(0.5)
print(images[:,mask,:,:].size())

#OR just treat all channels and batches together
mask2 = torch.randn(100*3).ge(0.5)
images=images.view(3*100, 224,224)

images = images[mask2,...]
print(images.size())
</code></pre>

jack_S: <p>Thanks. The second way works. However, the size is fixed on a given dim so I thought pytorch has methods to slice this.</p>

J_Johnson: <p>Not sure your use case, but have you considered just setting the values you don’t want to zero?</p>
<p>Perhaps you can share what your desired output should look like or your overall objective to help us narrow down a solution for you.</p>

jack_S: <p>Ok I’ll say clearer. Now I use a smaller tensor to discribe my disired output.</p>
<pre><code class="lang-auto">images = torch.linspace(0, 14, 15)
images = images.view(5, 3, 1, 1).expand(5, 3, 10, 10) 
#images.shape = [5, 3, 10, 10] and 
#elements in images[i, j] are all 3 * i + j
mask = torch.tensor([[True, False, False],
                     [False, True, False],
                     [False, False, True],
                     [False, True, False],
                     [True, False, False]])
</code></pre>
<p>after a pytorch-style slicing the output <code>output_tensor</code> should be</p>
<pre><code class="lang-auto">output_tensor = torch.tensor([[[[0] * 10] * 10], 
                              [[[4] * 10] * 10],
                              [[[8] * 10] * 10],
                              [[[10] * 10] * 10],
                              [[[12] * 10] * 10]])
#output_tensor.shape = [5, 1, 10, 10]
</code></pre>
<p>Or if the <code>mask</code> is</p>
<pre><code class="lang-auto">mask = torch.tensor([[True, True, False],
                     [False, True, True],
                     [True, False, True],
                     [True, True, False],
                     [True, False, True]])
</code></pre>
<p>after a pytorch-style slicing the output <code>output_tensor</code> should be</p>
<pre><code class="lang-auto">output_tensor = torch.tensor([[[[0] * 10] * 10, [[1] * 10] * 10], 
                              [[[4] * 10] * 10, [[5] * 10] * 10],
                              [[[6] * 10] * 10, [[8] * 10] * 10],
                              [[[9] * 10] * 10, [[10] * 10] * 10],
                              [[[12] * 10] * 10, [[14] * 10] * 10]])
#output_tensor.shape = [5, 2, 10, 10]
</code></pre>

J_Johnson: <p>Looks to me in both cases <code>dim == 1</code> is the same within each respective batch. You could just modify the 2nd option to resize that dim dynamically. Here is how that might work:</p>
<pre><code class="lang-auto">images=torch.randn(100, 3, 224, 224)

mask1 = torch.randint(high=3, size=(100,3)) &gt; 1 # for mask.size(1) == 1
mask2 = torch.randint(high=3, size=(100,3)) &gt; 0 # for mask.size(1) == 2

b, c, h, w = images.shape

images=images.view(3*100, 224,224)

images1 = images[mask1.view(-1),...].view(b, -1, h, w)
images2 = images[mask2.view(-1),...].view(b, -1, h, w)
</code></pre>
<p><em>currently not at a computer to test the above, so please don’t mind the typos</em></p>

jack_S: <p>Yes this is a great slicing method. I just think if there were a pytorch-style slicing way in one line (or say in one pair of brackets) then it would be more convenient.</p>