如果通道有任何大于 0 的值而不循环，则替换每个通道中的每个元素

Question

我有一个大小为 (4, 100, 56, 56) 的批次张量，其中一些通道具有特定值，而另一些通道只有全零。我想让通道中的每个元素都具有大于 0 的任何值，即 100，而如果它全为零，则应该在每个元素中设置为 1。知道如何在不循环的情况下实现这一点吗？

t = torch.zeros((4, 100, 56, 56))
t[:, 5, 15:20, 15:20] = 0.07
new_t = torch.ones((4, 100, 56, 56))

for b in range(t.size(0))
    for c in range(t.size(1)):
        if t[b, c, :,:].max() > 0:
             new_t[b, c, :, :] = 100

我上面的代码对于大批量和通道来说效率低下，并且由于new_t它会产生内存开销，有没有办法使用view()或类似的函数来实现这个？

Answer 1

您可以执行以下操作：

mask = torch.any(t.flatten(2, 3) > 0., dim=2)
t[mask] = 100.  # or t[mask] *= 100. for differentiability

如果通道有任何大于 0 的值而不循环，则替换每个通道中的每个元素

Replace every element in each channel if the channel has any value greater than 0 without looping

pytorch