arrays - Matlab:Arrayfun 与矩阵

我正在尝试使用 arrayfun 重写以下代码

A = ones(3,3,3)
for i = 1:3
    B(i) = trace(A(:,:,i));
end

我希望尝试过

f = @(x) trace(x)
B = arrayfun(f, A);

但这只是(如您所料)跟踪每个单独的 A(i,j,k) 而不是我想要的 A(:,:,i)。然后我尝试将 A{i}=ones(3,3) 声明为一个单元格并传递给 arrayfun 但这也不起作用。

如何在 Matlab 中向量化矩阵上的函数？

最佳答案

基于 bsxfun 的矢量化解决方案，ab(使用)如何定义 trace - sum of diagonal elements -

%// Get size of A
[m,n,r] = size(A)

%// Get indices of the diagonal elements for each 3D "slice" as columns of idx
idx = bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n) %//'

%// Thus, for your 3 x 3 x 3 case, idx would be -
%//idx =
%//      1    10    19
%//      5    14    23
%//      9    18    27
%// and these are the linear indices to the diagonal elements to each `3D` slide.

%//Index into A with idx and sum along columns to get each element of desired output
B = sum(A(idx),1)

如果您想使用不必要的额外变量来避免工作区困惑，请避免使用 idx -

B = sum(A(bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n)),1)

用于使用 GPU

如果您必须使用 GPUs ，您可以使用 gpuArray(A) 将它们声明为 gpuArrays，然后涉及 A 的后续工作将在 GPU 上完成，您将获得 gpuArray 的输出，您可以使用 gather(..) 作为 CPU 变量返回。

因此，完整的代码如下所示 -

[m,n,r] = size(A); %// Get size
gpu_A = gpuArray(A); %// copy data from CPU to GPU

%// Perform calculations on GPU
gpu_B = sum(gpu_A(bsxfun(@plus,[1:m+1:m*n]',[0:r-1]*m*n)),1); %//'

B = gather(gpu_B); %// get back output onto CPU

快速测试: 使用 GTX 750 Ti(我可以访问)，这似乎使我的循环代码速度提高了 3 倍。