从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