计算 3D 矩阵中的所有非零元素，但对于维度 2 中的每个平面

Question

我想执行与此功能相同的操作，但在内存中，我真的不知道该怎么做。

操作：输入是一个扫描，由叠加图像组成，产生一个维度（高度、宽度、图像数量），我想使用单个操作获得一个具有维度（图像数量）的向量，其中值将对于扫描中的每个图像都是非零元素。

我能做的最好的事情就是遍历每个图像。

def get_plane_where_tumor_is_thelargest_v2(brain_scan_label: torch.Tensor):
    slices_areas = torch.zeros(brain_scan_label.shape[-1])
    for i in range(brain_scan_label.shape[-1]):
        slices_areas[i] = brain_scan_label[:, :, i].count_nonzero()
        
    return slices_areas.argmax()

编辑：

我设法解决了这个问题，我没有看到一次将多个维度传递给 torch.count_nonzero() 的可能性，这样做就足够了：

brain_scan_label.count_nonzero(dim=(0, 1))

Answer 1

我会写这个作为答案，因为由于声誉点数不足，我还不能评论。

虽然我以前没有使用过 torch，但 count_nonzero() 函数似乎有可选的 dim 参数，它应该允许您指定要计算非零元素的张量轴。有关详细信息和一些示例，请参阅此链接和此其他链接。

在你的情况下，我的猜测是简单地写

brain_scan_label.count_nonzero(dim=2)

，这样你就可以避免循环。