Matlab中沿维度子集的多维数组的Argmax

Question

说，Y 是一个 7 维数组，我需要一种有效的方法来沿最后 3 维最大化它，这将在 GPU 上工作。 因此，我需要一个最大值为 Y 的 4 维数组和三个 4 维数组，这些数组在最后三个维度中具有这些值的索引。 我可以的

[Y7, X7] = max(Y , [], 7);
[Y6, X6] = max(Y7, [], 6);
[Y5, X5] = max(Y6, [], 5);

然后我已经找到了沿第 5 维 (X5) 的值 (Y5) 和索引。但我仍然需要第 6 维和第 7 维的索引。

Answer 1

这是一种方法。让N 表示要最大化的维度数。

1. 重塑 Y 以将最后一个 N 维度合并为一个。
2. 沿折叠尺寸最大化。这将 argmax 作为这些维度上的线性索引。
3. 将线性索引展开为N 子索引，每个维度一个。

以下代码适用于任意数量的维度（不一定是您的示例中的7 和3）。为了实现这一点，它通常处理 Y 的大小，并使用从元胞数组中获取的逗号分隔列表从 sub2ind 获取 N 输出。

Y = rand(2,3,2,3,2,3,2); % example 7-dimensional array
N = 3; % last dimensions along which to maximize
D = ndims(Y);
sz = size(Y);
[~, ind] = max(reshape(Y, [sz(1:D-N) prod(sz(D-N+1:end))]), [], D-N+1);
sub = cell(1,N);
[sub{:}] = ind2sub(sz(D-N+1:D), ind);

作为检查，运行上述代码后，观察例如Y(2,3,1,2,:)（为方便起见显示为行向量）：

>> reshape(Y(2,3,1,2,:), 1, [])
ans =
    0.5621    0.4352    0.3672    0.9011    0.0332    0.5044    0.3416    0.6996    0.0610    0.2638    0.5586    0.3766

最大值为0.9011，出现在4th 位置（其中“位置”沿N=3 折叠维度定义）。事实上，

>> ind(2,3,1,2)
ans =
     4
>> Y(2,3,1,2,ind(2,3,1,2))
ans =
    0.9011

或者，就N=3 子索引而言，

>> Y(2,3,1,2,sub{1}(2,3,1,2),sub{2}(2,3,1,2),sub{3}(2,3,1,2))
ans =
    0.9011