矢量化 pytorch 张量索引

Question

我有一批图像 img_batch，大小为 [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 个 dims 的遮罩求和来检查遮罩是否设置正确，并查看它是否符合您的目标百分比：

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

您无需循环即可以完全矢量化的方式轻松完成此操作：

创建噪声张量
Select 一个阈值并根据高于或低于该阈值 (prct0)
Element-wise 图像张量乘以噪声张量

我认为调用幂乘数向量 noise 有点令人困惑，因此我在本例中将该向量重命名为 power_vec： power_vec = 噪音

# 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)

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

矢量化 pytorch 张量索引

vectorize pytorch tensor indexing

indexing

pytorch

vectorization

tensor