向量化pytorch张量索引

Question

我有一批图像img_batch, size[8,3,32,32]，我想通过将随机选择的像素设置为零来操作每个图像。我可以在每个图像上使用for循环，但我不确定如何向量化它，所以我不是一次只处理一个图像。这是我使用循环的代码。

batch_size = 8
prct0 = 0.1
noise = torch.tensor([9, 14, 5, 7, 6, 14, 1, 3])
comb_img = []
for ind in range(batch_size):
img = img_batch[ind]
c, h, w = img.shape          
prct = 1 - (1 - prct0)**noise[ind].item()
idx = random.sample(range(h*w), int(prct*h*w)  )
img_noised = img.clone()
img_noised.view(c,1,-1)[:,0,idx] = 0 
comb_img.append(img_noised)
comb_img = torch.stack(comb_img) # output is comb_img [8,3,32,32]

我是pytorch的新手，如果你看到任何其他改进，请分享。

Answer 1

第一个注意事项:您需要使用噪点吗?如果你把所有的图像都当作相同的，并且没有不同的像素集来设置为0，这将会容易得多。

然而，你可以这样做，但你仍然需要一个小的for循环(在列表推导式中)。

#don't want RGB masking, want the whole pixel
rng = torch.rand(*img_batch[:,0:1].shape) 
#create binary mask
mask = torch.stack([rng[i] <= 1-(1-prct0)**noise[i] for i in range(batch_size)]) 
img_batch_masked = img_batch.clone()
#broadcast mask to 3 RGB channels
img_batch_masked[mask.tile([1,3,1,1])] = 0

你可以通过将最后3个暗点的掩码相加来检查掩码是否设置正确，并查看它是否与目标百分比匹配:

In [5]:     print(mask.sum([1,2,3])/(mask.shape[2] * mask.shape[3]))
tensor([0.6058, 0.7716, 0.4195, 0.5162, 0.4739, 0.7702, 0.1012, 0.2684])
In [6]:     print(1-(1-prct0)**noise)
tensor([0.6126, 0.7712, 0.4095, 0.5217, 0.4686, 0.7712, 0.1000, 0.2710])

Answer 2

您可以轻松地以完全矢量化的方式完成此操作，而无需循环:

1. 创建噪声张量
2. 选择一个阈值，并根据高于或低于该阈值(prct0)将噪声张量四舍五入到0或1
3. 逐元素的图像张量乘以噪声张量

我认为将功率乘法器的向量称为noise有点令人困惑，所以我在本例中将该向量重命名为power_vec:Power_vec = noise

# create random noise - note one channel rather than 3 color channels
rand_noise = torch.rand(8,1,32,32)
noise = torch.pow(rand_noise,power_vec) # these tensors are broadcastable

# "round" noise based on threshold  
z = torch.zeros(noise.shape)
o = torch.ones(noise.shape)
noise_rounded = torch.where(noise>prct0,o,z) 
# apply noise mask to each color channel
output = img_batch * noise_rounded.expand(8,3,32,32)    

为简单起见，此解决方案使用原始的批处理大小和图像大小，但可以简单地扩展到任何图像和批处理大小的输入。