如何在PyTorch / Numpy中使用不同尺寸的蒙版？

Question

我有一个大小为mask的torch.Size([20, 1, 199])和张量，reconstruct_output和inputs都大小为torch.Size([20, 1, 161, 199])。

我想将reconstruct_output设置为inputs，其中mask为0。我尝试过：

reconstruct_output[mask == 0] = inputs[mask == 0]

但出现错误：

IndexError: The shape of the mask [20, 1, 199] at index 2 does not match the shape of the indexed tensor [20, 1, 161, 199] at index 2

Answer 1

我们可以在这里使用advanced indexing。为了获得我们要用于索引advanced indexing和reconstruct_output的索引数组，我们需要沿inputs的轴进行索引。为此，我们可以使用m==0，并使用结果索引将np.where更新为：