查找具有最多 1 的 PyTorch 张量的列索引

Question

我有一个PyTorch 张量a，形状如下：

import torch
a = torch.tensor([[[1., 0., 0., 0.]],
        [[0., 1., 0., 0.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]]])

张量a的每一行有4个元素，1和0。假设我相应地索引了这个张量的行和列。例如，第 0 行（最上一行）中的条目是 [[1., 0., 0., 0.]]，而第 3 列（最右列）中的条目是 [[0., 0., 0., 1., 0., 1., 0.]]。

从给定的张量中，我想确定1. 出现频率最高的列的索引。例如，对于张量a，此类列的索引将为 0。如果 1. 的数量存在关联，我仍然希望获取所有关联列索引。

如何在 Python 上轻松完成这项任务？

谢谢，

Answer 1

如果你的矩阵只包含0和1，你可以对每一列的元素求和，然后搜索最大的和：

import numpy as np

% sum over columns
sumsi = torch.sum(a, dim=1)

% find where maximum
col_idx = np.where(sumsi==np.max(sumsi))