I have a batch of images (4d tensor/array with dimensions "batchsize x channels x height x width" and I would like to draw horizontal bars of zeros of size s on each image, but across different rows for each image. I can do this trivially with a for loop, but I haven't been able to figure out a vectorized implementation.
Ideally I would generate a 1-D tensor r of "batchsize" random starting points, and do something like
t[:,:,r:r+s,:] = 0. If I try this I get TypeError: only integer scalar arrays can be converted to a scalar index
If I do a toy example and just try to pull out two different sections of a batch with only two images, doing something like t[:,:,torch.tensor(([1,2],[2,3])),:] I get back a 5D tensor because it is pulling both of those sections from both images in the batch. How do I grab those different sections but only one for each image? In this case if the input were 2xCxHxW I would want 2xCx2xW where the first item corresponds to rows 1 and 2 of the first image, and the second item corresponds to rows 2 and 3 of the second image. Thank you.
You can use this function which will create a mask where you can perform operations across the y or x axis by their index. You can do this by arranging the x values of the index to be set to their y index.
bsg = sgs.data
device = sgs.device
bs, _, x, y = bsg.shape
max_y = y-size-1
rs = torch.randint(0, max_y, (bs,1), device=device)
m = torch.arange(y,device=device).repeat(bs, x)
gpumask = ((m < rs) | (m > (rs+size))).view(bs, 1, x, -1)
gpumask*bsg