在 Pytorch 中从 4D 张量中查找零索引

Question

如果你正在使用 Pytorch，我有一个有趣的问题要问你。

我有一批形状为 (3, 224, 224) 的图像。因此，如果我的批量大小为 64，假设我的最终张量的形状为 (64, 3, 224, 224)。

现在，问题来了。假设这批中的一些图像只填充了零。找出哪些批次索引只有零的最快方法是什么？

我不想为此创建 for 循环，因为它很慢。

感谢您的回答。

Answer 1

一种更便宜的方法是假设只有空图像的总和 = 0，我认为这是非常合理的

import torch
t = torch.rand(64,3,224,224)
t[10] = 0
s = t.view(64, -1).sum(dim = -1)
zero_index = (s==0).nonzero().item() # 10